diff --git a/CMakeLists.txt b/CMakeLists.txt
index deeb99c20..803698f10 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -36,31 +36,24 @@ macro(torch_mlir_add_llvm_external_project name identifier location)
   set(LLVM_EXTERNAL_PROJECTS ${LLVM_EXTERNAL_PROJECTS} CACHE STRING "" FORCE)
 endmacro()
 
+#-------------------------------------------------------------------------------
+# Project options
+#-------------------------------------------------------------------------------
+
+option(TORCH_MLIR_ENABLE_REFBACKEND "Enable reference backend" ON)
+if(TORCH_MLIR_ENABLE_REFBACKEND)
+  add_definitions(-DTORCH_MLIR_ENABLE_REFBACKEND)
+endif()
+
 option(TORCH_MLIR_ENABLE_STABLEHLO "Add stablehlo dialect" ON)
 if(TORCH_MLIR_ENABLE_STABLEHLO)
   add_definitions(-DTORCH_MLIR_ENABLE_STABLEHLO)
 endif()
 
-option(TORCH_MLIR_ENABLE_JIT_IR_IMPORTER "Enables JIT IR Importer" ON)
-option(TORCH_MLIR_ENABLE_LTC "Enables LTC backend" OFF)
-option(TORCH_MLIR_ENABLE_ONLY_MLIR_PYTHON_BINDINGS "Build Torch dialect MLIR Python bindings but neither JIT IR Importer nor LTC backend" OFF)
-if(TORCH_MLIR_ENABLE_ONLY_MLIR_PYTHON_BINDINGS)
-  set(TORCH_MLIR_ENABLE_JIT_IR_IMPORTER OFF)
-  set(TORCH_MLIR_ENABLE_LTC OFF)
-endif()
-
-if(TORCH_MLIR_ENABLE_LTC)
-    set(ENV{TORCH_MLIR_ENABLE_LTC} 1)
-    message(STATUS "LTC Backend build is enabled")
-else()
-    set(ENV{TORCH_MLIR_ENABLE_LTC} 0)
-    message(STATUS "LTC Backend build is disabled")
-endif()
-
-torch_mlir_add_llvm_external_project(
-  torch-mlir-dialects
-  TORCH_MLIR_DIALECTS
-  ${CMAKE_CURRENT_SOURCE_DIR}/externals/llvm-external-projects/torch-mlir-dialects)
+# torch_mlir_add_llvm_external_project(
+#   torch-mlir-dialects
+#   TORCH_MLIR_DIALECTS
+#   ${CMAKE_CURRENT_SOURCE_DIR}/externals/llvm-external-projects/torch-mlir-dialects)
 
 option(TORCH_MLIR_OUT_OF_TREE_BUILD "Specifies an out of tree build" OFF)
 
@@ -99,9 +92,9 @@ if(CMAKE_SOURCE_DIR STREQUAL CMAKE_CURRENT_SOURCE_DIR OR TORCH_MLIR_OUT_OF_TREE_
 
   # Don't try to compile the python extensions at the moment. We need
   # to import lots of dependencies from AddMLIRPython to make this work.
-  set(MLIR_ENABLE_BINDINGS_PYTHON 1)
+  set(MLIR_ENABLE_BINDINGS_PYTHON ON)
 
-  set(TORCH-MLIR_BUILT_STANDALONE 1)
+  set(TORCH-MLIR_BUILT_STANDALONE ON)
   set(BACKEND_PACKAGE_STRING "LLVM ${LLVM_PACKAGE_VERSION}")
   add_subdirectory(externals/llvm-external-projects/torch-mlir-dialects)
 else()
@@ -181,10 +174,9 @@ if(MLIR_ENABLE_BINDINGS_PYTHON)
   if(NOT TORCH_MLIR_PYTHON_PACKAGES_DIR)
     set(TORCH_MLIR_PYTHON_PACKAGES_DIR "${CMAKE_CURRENT_BINARY_DIR}/python_packages")
   endif()
-  add_dependencies(check-torch-mlir-all
-  check-torch-mlir-python
-  )
-  add_subdirectory(python)
+  #add_subdirectory(python)
+  # TODO: Remove
+  add_custom_target(TorchMLIRPythonModules)
 else()
   add_custom_target(TorchMLIRPythonModules)
 endif()
diff --git a/include/torch-mlir/InitAll.h b/include/torch-mlir/InitAll.h
index ecced42cd..42eb3c6a1 100644
--- a/include/torch-mlir/InitAll.h
+++ b/include/torch-mlir/InitAll.h
@@ -15,7 +15,13 @@
 namespace mlir {
 namespace torch {
 
+// Registers all dialects that this project produces and any dependencies.
 void registerAllDialects(mlir::DialectRegistry &registry);
+
+// Registers dialects that may be needed to parse torch-mlir inputs and
+// test cases.
+void registerOptionalInputDialects(mlir::DialectRegistry &registry);
+
 void registerAllPasses();
 
 } // namespace torch
diff --git a/lib/CAPI/CMakeLists.txt b/lib/CAPI/CMakeLists.txt
index d71796ae8..927fb7d3c 100644
--- a/lib/CAPI/CMakeLists.txt
+++ b/lib/CAPI/CMakeLists.txt
@@ -9,8 +9,6 @@ add_mlir_public_c_api_library(TorchMLIRCAPI
   ${PROJECT_SOURCE_DIR}/include/torch-mlir-c/
 
   ENABLE_AGGREGATION
-  LINK_COMPONENTS
-  Core
 
   LINK_LIBS PUBLIC
   MLIRIR
diff --git a/lib/CMakeLists.txt b/lib/CMakeLists.txt
index 03123d2ed..8956066b8 100644
--- a/lib/CMakeLists.txt
+++ b/lib/CMakeLists.txt
@@ -1,14 +1,18 @@
 add_subdirectory(CAPI)
 add_subdirectory(Conversion)
 add_subdirectory(Dialect)
-add_subdirectory(RefBackend)
 
-get_property(extension_libs GLOBAL PROPERTY MLIR_EXTENSION_LIBS)
 set(LinkedLibs
+  MLIRComplexDialect
   MLIRFuncDialect
+  MLIRFuncInlinerExtension
   MLIRIR
+  MLIRMLProgramDialect
+  MLIRMemRefDialect
+  MLIRSCFDialect
+  MLIRTensorDialect
+  MLIRTosaDialect
   MLIRSupport
-  ${extension_libs}
 
   TorchMLIRTorchPasses
   TorchMLIRTorchConversionDialect
@@ -20,15 +24,16 @@ set(LinkedLibs
   TorchMLIRTMTensorDialect
 
   TorchMLIRConversionPasses
-  TorchMLIRRefBackend
 )
 
+if(TORCH_MLIR_ENABLE_REFBACKEND)
+  add_subdirectory(RefBackend)
+  list(APPEND LinkedLibs TorchMLIRRefBackend)
+endif()
+
 add_mlir_library(TorchMLIRInitAll
   InitAll.cpp
 
-  LINK_COMPONENTS
-  Core
-
   LINK_LIBS PUBLIC
   ${LinkedLibs}
 )
diff --git a/lib/Conversion/CMakeLists.txt b/lib/Conversion/CMakeLists.txt
index d72563b1e..f26b4d6e8 100644
--- a/lib/Conversion/CMakeLists.txt
+++ b/lib/Conversion/CMakeLists.txt
@@ -27,10 +27,6 @@ add_mlir_library(TorchMLIRConversionPasses
   DEPENDS
   TorchMLIRConversionPassIncGen
 
-  LINK_COMPONENTS
-  Core
-
   LINK_LIBS PUBLIC
   ${linked_libs}
-  #${torch_mlir_conversion_libs}
 )
diff --git a/lib/Conversion/TorchConversionToMLProgram/CMakeLists.txt b/lib/Conversion/TorchConversionToMLProgram/CMakeLists.txt
index b89ffbb43..3546861ce 100644
--- a/lib/Conversion/TorchConversionToMLProgram/CMakeLists.txt
+++ b/lib/Conversion/TorchConversionToMLProgram/CMakeLists.txt
@@ -7,9 +7,6 @@ add_mlir_conversion_library(TorchMLIRTorchConversionToMLProgram
   DEPENDS
   TorchMLIRConversionPassIncGen
 
-  LINK_COMPONENTS
-  Core
-
   LINK_LIBS PUBLIC
   MLIRIR
   MLIRLinalgDialect
diff --git a/lib/Conversion/TorchToArith/CMakeLists.txt b/lib/Conversion/TorchToArith/CMakeLists.txt
index 4524c3b07..b345f3c96 100644
--- a/lib/Conversion/TorchToArith/CMakeLists.txt
+++ b/lib/Conversion/TorchToArith/CMakeLists.txt
@@ -7,9 +7,6 @@ add_mlir_conversion_library(TorchMLIRTorchToArith
   DEPENDS
   TorchMLIRConversionPassIncGen
 
-  LINK_COMPONENTS
-  Core
-
   LINK_LIBS PUBLIC
   MLIRIR
   MLIRPass
diff --git a/lib/Conversion/TorchToLinalg/CMakeLists.txt b/lib/Conversion/TorchToLinalg/CMakeLists.txt
index ece929597..fbf66556e 100644
--- a/lib/Conversion/TorchToLinalg/CMakeLists.txt
+++ b/lib/Conversion/TorchToLinalg/CMakeLists.txt
@@ -17,9 +17,6 @@ add_mlir_conversion_library(TorchMLIRTorchToLinalg
   DEPENDS
   TorchMLIRConversionPassIncGen
 
-  LINK_COMPONENTS
-  Core
-
   LINK_LIBS PUBLIC
   MLIRIR
   MLIRPass
diff --git a/lib/Conversion/TorchToSCF/CMakeLists.txt b/lib/Conversion/TorchToSCF/CMakeLists.txt
index 06a888b3e..7a6d21cd6 100644
--- a/lib/Conversion/TorchToSCF/CMakeLists.txt
+++ b/lib/Conversion/TorchToSCF/CMakeLists.txt
@@ -7,9 +7,6 @@ add_mlir_conversion_library(TorchMLIRTorchToSCF
   DEPENDS
   TorchMLIRConversionPassIncGen
 
-  LINK_COMPONENTS
-  Core
-
   LINK_LIBS PUBLIC
   MLIRIR
   MLIRPass
diff --git a/lib/Conversion/TorchToStablehlo/CMakeLists.txt b/lib/Conversion/TorchToStablehlo/CMakeLists.txt
index 0f9b8faba..07ef1e2ea 100644
--- a/lib/Conversion/TorchToStablehlo/CMakeLists.txt
+++ b/lib/Conversion/TorchToStablehlo/CMakeLists.txt
@@ -14,9 +14,6 @@ add_mlir_conversion_library(TorchMLIRTorchToStablehlo
   DEPENDS
   TorchMLIRConversionPassIncGen
 
-  LINK_COMPONENTS
-  Core
-
   LINK_LIBS PUBLIC
   MLIRIR
   MLIRPass
diff --git a/lib/Conversion/TorchToTMTensor/CMakeLists.txt b/lib/Conversion/TorchToTMTensor/CMakeLists.txt
index d05d8277c..3f0cf9de9 100644
--- a/lib/Conversion/TorchToTMTensor/CMakeLists.txt
+++ b/lib/Conversion/TorchToTMTensor/CMakeLists.txt
@@ -7,9 +7,6 @@ TorchToTMTensor.cpp
   DEPENDS
   TorchMLIRConversionPassIncGen
 
-  LINK_COMPONENTS
-  Core
-
   LINK_LIBS PUBLIC
   MLIRIR
   MLIRPass
diff --git a/lib/Conversion/TorchToTosa/CMakeLists.txt b/lib/Conversion/TorchToTosa/CMakeLists.txt
index 909ee3bcb..24f349902 100644
--- a/lib/Conversion/TorchToTosa/CMakeLists.txt
+++ b/lib/Conversion/TorchToTosa/CMakeLists.txt
@@ -9,9 +9,6 @@ add_mlir_conversion_library(TorchMLIRTorchToTosa
   DEPENDS
   TorchMLIRConversionPassIncGen
 
-  LINK_COMPONENTS
-  Core
-
   LINK_LIBS PUBLIC
   MLIRIR
   MLIRPass
diff --git a/lib/Dialect/Torch/IR/CMakeLists.txt b/lib/Dialect/Torch/IR/CMakeLists.txt
index 00210e4fd..58b7cdd0b 100644
--- a/lib/Dialect/Torch/IR/CMakeLists.txt
+++ b/lib/Dialect/Torch/IR/CMakeLists.txt
@@ -12,9 +12,6 @@ add_mlir_library(TorchMLIRTorchDialect
   MLIRTorchOpsIncGen
   MLIRTorchTypesIncGen
 
-  LINK_COMPONENTS
-  Core
-
   LINK_LIBS PUBLIC
   MLIRBytecodeOpInterface
   MLIRBytecodeReader
diff --git a/lib/Dialect/Torch/Transforms/CMakeLists.txt b/lib/Dialect/Torch/Transforms/CMakeLists.txt
index 9c8bcda94..0f7621ff0 100644
--- a/lib/Dialect/Torch/Transforms/CMakeLists.txt
+++ b/lib/Dialect/Torch/Transforms/CMakeLists.txt
@@ -26,9 +26,6 @@ add_mlir_library(TorchMLIRTorchPasses
   DEPENDS
   TorchMLIRTorchPassIncGen
 
-  LINK_COMPONENTS
-  Core
-
   LINK_LIBS PUBLIC
   MLIRIR
   MLIRPass
diff --git a/lib/Dialect/TorchConversion/IR/CMakeLists.txt b/lib/Dialect/TorchConversion/IR/CMakeLists.txt
index 38f6f3268..cfe7006f7 100644
--- a/lib/Dialect/TorchConversion/IR/CMakeLists.txt
+++ b/lib/Dialect/TorchConversion/IR/CMakeLists.txt
@@ -9,9 +9,6 @@ add_mlir_dialect_library(TorchMLIRTorchConversionDialect
   MLIRTorchConversionOpsIncGen
   MLIRTorchTypesIncGen
 
-  LINK_COMPONENTS
-  Core
-
   LINK_LIBS PUBLIC
   MLIRIR
   MLIRSupport
diff --git a/lib/Dialect/TorchConversion/Transforms/CMakeLists.txt b/lib/Dialect/TorchConversion/Transforms/CMakeLists.txt
index a286d5bbd..5858e29a4 100644
--- a/lib/Dialect/TorchConversion/Transforms/CMakeLists.txt
+++ b/lib/Dialect/TorchConversion/Transforms/CMakeLists.txt
@@ -39,9 +39,6 @@ add_mlir_library(TorchMLIRTorchConversionPasses
   DEPENDS
   TorchMLIRTorchConversionPassIncGen
 
-  LINK_COMPONENTS
-  Core
-
   LINK_LIBS PUBLIC
   ${LinkedLibs}
 )
diff --git a/lib/InitAll.cpp b/lib/InitAll.cpp
index 1d67bdfe2..0be0ec8ba 100644
--- a/lib/InitAll.cpp
+++ b/lib/InitAll.cpp
@@ -9,8 +9,15 @@
 
 #include "torch-mlir/InitAll.h"
 
+#include "mlir/Dialect/Complex/IR/Complex.h"
 #include "mlir/Dialect/Func/Extensions/InlinerExtension.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/Dialect/Linalg/IR/Linalg.h"
+#include "mlir/Dialect/MLProgram/IR/MLProgram.h"
+#include "mlir/Dialect/MemRef/IR/MemRef.h"
+#include "mlir/Dialect/SCF/IR/SCF.h"
+#include "mlir/Dialect/Tensor/IR/Tensor.h"
+#include "mlir/Dialect/Tosa/IR/TosaOps.h"
 #include "mlir/IR/Dialect.h"
 #include "torch-mlir-dialects/Dialect/TMTensor/IR/TMTensorDialect.h"
 #include "torch-mlir-dialects/Dialect/TMTensor/Transforms/Passes.h"
@@ -29,11 +36,22 @@ void mlir::torch::registerAllDialects(mlir::DialectRegistry &registry) {
   mlir::func::registerInlinerExtension(registry);
 }
 
+// TODO: Break this up when backends are separated.
+void mlir::torch::registerOptionalInputDialects(
+    mlir::DialectRegistry &registry) {
+  registry.insert<complex::ComplexDialect, linalg::LinalgDialect,
+                  memref::MemRefDialect, ml_program::MLProgramDialect,
+                  scf::SCFDialect, tensor::TensorDialect, tosa::TosaDialect>();
+}
+
 void mlir::torch::registerAllPasses() {
   mlir::torch::registerTorchPasses();
   mlir::torch::registerTorchConversionPasses();
 
   mlir::torch::registerConversionPasses();
-  mlir::torch::RefBackend::registerRefBackendPasses();
   mlir::torch::TMTensor::registerPasses();
+
+#ifdef TORCH_MLIR_ENABLE_REFBACKEND
+  mlir::torch::RefBackend::registerRefBackendPasses();
+#endif
 }
diff --git a/test/CMakeLists.txt b/test/CMakeLists.txt
index 51407b488..8f4cf9f3d 100644
--- a/test/CMakeLists.txt
+++ b/test/CMakeLists.txt
@@ -1,6 +1,6 @@
 llvm_canonicalize_cmake_booleans(
   MLIR_ENABLE_BINDINGS_PYTHON
-  TORCH_MLIR_ENABLE_JIT_IR_IMPORTER
+  TORCH_MLIR_ENABLE_REFBACKEND
   TORCH_MLIR_ENABLE_STABLEHLO
 )
 
diff --git a/test/RefBackend/lit.local.cfg b/test/RefBackend/lit.local.cfg
new file mode 100644
index 000000000..344ff6e4c
--- /dev/null
+++ b/test/RefBackend/lit.local.cfg
@@ -0,0 +1,2 @@
+if not config.torch_mlir_enable_refbackend:
+    config.unsupported = True
diff --git a/test/lit.site.cfg.py.in b/test/lit.site.cfg.py.in
index f7ac86747..1be54aaf6 100644
--- a/test/lit.site.cfg.py.in
+++ b/test/lit.site.cfg.py.in
@@ -5,6 +5,7 @@ import sys
 config.enable_bindings_python = @MLIR_ENABLE_BINDINGS_PYTHON@
 config.torch_mlir_obj_root = "@TORCH_MLIR_BINARY_DIR@"
 config.torch_mlir_python_packages_dir = "@TORCH_MLIR_PYTHON_PACKAGES_DIR@"
+config.torch_mlir_enable_refbackend = @TORCH_MLIR_ENABLE_REFBACKEND@
 config.host_os = "@HOST_OS@"
 config.llvm_src_root = "@LLVM_SOURCE_DIR@"
 config.llvm_obj_root = "@LLVM_BINARY_DIR@"
@@ -16,7 +17,6 @@ config.llvm_shlib_ext = "@SHLIBEXT@"
 config.llvm_exe_ext = "@EXEEXT@"
 config.lit_tools_dir = "@LLVM_LIT_TOOLS_DIR@"
 config.python_executable = "@Python3_EXECUTABLE@"
-config.enable_jit_ir_importer = @TORCH_MLIR_ENABLE_JIT_IR_IMPORTER@
 config.enable_stablehlo = @TORCH_MLIR_ENABLE_STABLEHLO@
 
 import lit.llvm
diff --git a/tools/torch-mlir-lsp-server/CMakeLists.txt b/tools/torch-mlir-lsp-server/CMakeLists.txt
index d53519c8a..55dbd3e41 100644
--- a/tools/torch-mlir-lsp-server/CMakeLists.txt
+++ b/tools/torch-mlir-lsp-server/CMakeLists.txt
@@ -7,18 +7,18 @@ EXPORT TorchMLIRTargets
 RUNTIME DESTINATION ${LLVM_TOOLS_INSTALL_DIR}
 COMPONENT torch-mlir-lsp-server)
 
-get_property(dialect_libs GLOBAL PROPERTY MLIR_DIALECT_LIBS)
-get_property(conversion_libs GLOBAL PROPERTY MLIR_CONVERSION_LIBS)
-get_property(extension_libs GLOBAL PROPERTY MLIR_EXTENSION_LIBS)
+# get_property(dialect_libs GLOBAL PROPERTY MLIR_DIALECT_LIBS)
+# get_property(conversion_libs GLOBAL PROPERTY MLIR_CONVERSION_LIBS)
+# get_property(extension_libs GLOBAL PROPERTY MLIR_EXTENSION_LIBS)
 
 target_link_libraries(torch-mlir-lsp-server PRIVATE
   MLIRLspServerLib
   TorchMLIRInitAll
 
-  # TODO: Remove these in favor of interface deps.
-  ${dialect_libs}
-  ${conversion_libs}
-  ${extension_libs}
+  # # TODO: Remove these in favor of interface deps.
+  # ${dialect_libs}
+  # ${conversion_libs}
+  # ${extension_libs}
 )
 
 mlir_check_all_link_libraries(torch-mlir-lsp-server)
diff --git a/tools/torch-mlir-lsp-server/torch-mlir-lsp-server.cpp b/tools/torch-mlir-lsp-server/torch-mlir-lsp-server.cpp
index a6d88a355..4b66b2797 100644
--- a/tools/torch-mlir-lsp-server/torch-mlir-lsp-server.cpp
+++ b/tools/torch-mlir-lsp-server/torch-mlir-lsp-server.cpp
@@ -9,8 +9,6 @@
 
 #include "mlir/IR/Dialect.h"
 #include "mlir/IR/MLIRContext.h"
-#include "mlir/InitAllDialects.h"
-#include "mlir/InitAllExtensions.h"
 #include "mlir/Tools/mlir-lsp-server/MlirLspServerMain.h"
 #include "torch-mlir/InitAll.h"
 
@@ -18,8 +16,7 @@ using namespace mlir;
 
 int main(int argc, char **argv) {
   DialectRegistry registry;
-  registerAllDialects(registry);
-  registerAllExtensions(registry);
   mlir::torch::registerAllDialects(registry);
+  mlir::torch::registerOptionalInputDialects(registry);
   return failed(MlirLspServerMain(argc, argv, registry));
 }
diff --git a/tools/torch-mlir-opt/CMakeLists.txt b/tools/torch-mlir-opt/CMakeLists.txt
index 94c547d0e..5e2a81e04 100644
--- a/tools/torch-mlir-opt/CMakeLists.txt
+++ b/tools/torch-mlir-opt/CMakeLists.txt
@@ -5,10 +5,6 @@ EXPORT TorchMLIRTargets
 RUNTIME DESTINATION ${LLVM_TOOLS_INSTALL_DIR}
 COMPONENT torch-mlir-opt)
 
-get_property(dialect_libs GLOBAL PROPERTY MLIR_DIALECT_LIBS)
-get_property(conversion_libs GLOBAL PROPERTY MLIR_CONVERSION_LIBS)
-get_property(extension_libs GLOBAL PROPERTY MLIR_EXTENSION_LIBS)
-
 set(dependency_libraries)
 if(TORCH_MLIR_ENABLE_STABLEHLO)
   list(APPEND dependency_libraries StablehloRegister)
@@ -16,11 +12,9 @@ endif()
 
 target_link_libraries(torch-mlir-opt PRIVATE
   MLIROptLib
+  MLIRTransforms
   TorchMLIRInitAll
   TorchMLIRTorchDialect
   TorchMLIRTorchPasses
-  ${dialect_libs}
-  ${conversion_libs}
-  ${extension_libs}
   ${dependency_libraries}
 )
diff --git a/tools/torch-mlir-opt/torch-mlir-opt.cpp b/tools/torch-mlir-opt/torch-mlir-opt.cpp
index fa6a41a70..2750ee2b7 100644
--- a/tools/torch-mlir-opt/torch-mlir-opt.cpp
+++ b/tools/torch-mlir-opt/torch-mlir-opt.cpp
@@ -7,10 +7,10 @@
 //
 //===----------------------------------------------------------------------===//
 
-#include "mlir/InitAllDialects.h"
-#include "mlir/InitAllExtensions.h"
-#include "mlir/InitAllPasses.h"
+#include "mlir/Dialect/SCF/IR/SCF.h"
+#include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/Tools/mlir-opt/MlirOptMain.h"
+#include "mlir/Transforms/Passes.h"
 #include "torch-mlir/InitAll.h"
 
 #ifdef TORCH_MLIR_ENABLE_STABLEHLO
@@ -20,14 +20,23 @@
 using namespace mlir;
 
 int main(int argc, char **argv) {
-  registerAllPasses();
   mlir::torch::registerAllPasses();
 
+  // Core Transforms
+  registerCanonicalizerPass();
+  registerCSEPass();
+  registerInlinerPass();
+  registerLocationSnapshotPass();
+  registerLoopInvariantCodeMotionPass();
+  registerPrintOpStatsPass();
+  registerViewOpGraphPass();
+  registerStripDebugInfoPass();
+  registerSymbolDCEPass();
+
   DialectRegistry registry;
-  registerAllDialects(registry);
-  registerAllExtensions(registry);
   mlir::torch::registerAllDialects(registry);
-  
+  mlir::torch::registerOptionalInputDialects(registry);
+
 #ifdef TORCH_MLIR_ENABLE_STABLEHLO
   mlir::stablehlo::registerAllDialects(registry);
 #endif