Populate dependent dialects in auto input conversion pipeline (#13866) Because we don't statically know the exact pipeline, add dialects from all pipelines. Add a missing dialect to the MHLO pipeline. Fixes: https://github.com/openxla/iree/issues/13860

diff --git a/compiler/src/iree/compiler/InputConversion/Common/AutoInputConversionPipeline.cpp b/compiler/src/iree/compiler/InputConversion/Common/AutoInputConversionPipeline.cpp
index ea75d1f..c7233ad 100644
--- a/compiler/src/iree/compiler/InputConversion/Common/AutoInputConversionPipeline.cpp
+++ b/compiler/src/iree/compiler/InputConversion/Common/AutoInputConversionPipeline.cpp

@@ -27,10 +27,11 @@
 #endif  // IREE_HAVE_TORCH_INPUT
 
 namespace mlir::iree_compiler {
-
+namespace {
 struct AutoInputConversionPipelinePass
     : public AutoInputConversionPipelineBase<AutoInputConversionPipelinePass> {
   void runOnOperation() override;
+  void getDependentDialects(DialectRegistry& registry) const override;
 };
 
 // All the features seen that should be handled during input conversion.
@@ -166,6 +167,39 @@
   }
 }
 
+void AutoInputConversionPipelinePass::getDependentDialects(
+    DialectRegistry& registry) const {
+  // Register dialects from all possible pipelines, as we do not statically know
+  // which pipeline will be selected, while dialect registration happens before
+  // we run any detection on the input.
+  //
+  // TODO(kuhar): Find a better registration mechanism so that we do not have to
+  // build pipelines just to query dialects and discard them immediately after.
+  auto appendPipelineDialects =
+      [&registry](function_ref<void(OpPassManager&)> buildFn) {
+        OpPassManager pm;
+        buildFn(pm);
+        pm.getDependentDialects(registry);
+      };
+
+#ifdef IREE_HAVE_MHLO_INPUT
+  appendPipelineDialects(MHLO::buildMHLOInputConversionPassPipeline);
+  appendPipelineDialects(MHLO::buildXLAInputConversionPassPipeline);
+#endif  // IREE_HAVE_MHLO_INPUT
+
+#ifdef IREE_HAVE_TOSA_INPUT
+  appendPipelineDialects(buildTOSAInputConversionPassPipeline);
+#endif  // IREE_HAVE_TOSA_INPUT
+
+#ifdef IREE_HAVE_TORCH_INPUT
+  appendPipelineDialects([](OpPassManager& pm) {
+    pm.addNestedPass<func::FuncOp>(
+        TMTensor::createConvertTMTensorToLinalgExtPass());
+  });
+#endif  // IREE_HAVE_TORCH_INPUT
+}
+}  // namespace
+
 std::unique_ptr<OperationPass<ModuleOp>>
 createAutoInputConversionPipelinePass() {
   return std::make_unique<AutoInputConversionPipelinePass>();

diff --git a/compiler/src/iree/compiler/InputConversion/Common/test/BUILD.bazel b/compiler/src/iree/compiler/InputConversion/Common/test/BUILD.bazel
index 9a5367e..a99b55e 100644
--- a/compiler/src/iree/compiler/InputConversion/Common/test/BUILD.bazel
+++ b/compiler/src/iree/compiler/InputConversion/Common/test/BUILD.bazel

@@ -18,6 +18,7 @@
     name = "lit",
     srcs = enforce_glob(
         [
+            "auto_input_conversion_pipeline.mlir",
             "import_ml_program.mlir",
             "iree_import_public.mlir",
             "linalg_quantized_conv_to_conv.mlir",

diff --git a/compiler/src/iree/compiler/InputConversion/Common/test/CMakeLists.txt b/compiler/src/iree/compiler/InputConversion/Common/test/CMakeLists.txt
index f10539a..6c5ee1f 100644
--- a/compiler/src/iree/compiler/InputConversion/Common/test/CMakeLists.txt
+++ b/compiler/src/iree/compiler/InputConversion/Common/test/CMakeLists.txt

@@ -14,6 +14,7 @@
   NAME
     lit
   SRCS
+    "auto_input_conversion_pipeline.mlir"
     "import_ml_program.mlir"
     "iree_import_public.mlir"
     "linalg_quantized_conv_to_conv.mlir"

diff --git a/compiler/src/iree/compiler/InputConversion/Common/test/auto_input_conversion_pipeline.mlir b/compiler/src/iree/compiler/InputConversion/Common/test/auto_input_conversion_pipeline.mlir
new file mode 100644
index 0000000..1e8e90d
--- /dev/null
+++ b/compiler/src/iree/compiler/InputConversion/Common/test/auto_input_conversion_pipeline.mlir

@@ -0,0 +1,10 @@
+// RUN: iree-opt --iree-auto-input-conversion %s | FileCheck %s
+
+// Check that the input conversion pipeline handles a simple input and does not crash.
+
+// CHECK-LABEL: func.func @simple_add
+// CHECK:  arith.addi
+func.func @simple_add(%arg0: tensor<2x2xi32>, %arg1: tensor<2x2xi32>) -> tensor<2x2xi32> {
+  %0 = "mhlo.add"(%arg0, %arg1) : (tensor<2x2xi32>, tensor<2x2xi32>) -> tensor<2x2xi32>
+  return %0 : tensor<2x2xi32>
+}

diff --git a/compiler/src/iree/compiler/InputConversion/MHLO/MHLOToLinalgOnTensors.cpp b/compiler/src/iree/compiler/InputConversion/MHLO/MHLOToLinalgOnTensors.cpp
index 5363f6f..3992456 100644
--- a/compiler/src/iree/compiler/InputConversion/MHLO/MHLOToLinalgOnTensors.cpp
+++ b/compiler/src/iree/compiler/InputConversion/MHLO/MHLOToLinalgOnTensors.cpp

@@ -448,10 +448,10 @@
     : public ConvertMHLOToLinalgOnTensorsBase<
           ConvertMHLOToLinalgOnTensorsPass> {
   void getDependentDialects(DialectRegistry &registry) const override {
-    registry.insert<IREE::Flow::FlowDialect, IREE::Util::UtilDialect,
-                    linalg::LinalgDialect, mhlo::MhloDialect,
-                    shape::ShapeDialect, math::MathDialect,
-                    memref::MemRefDialect, complex::ComplexDialect>();
+    registry.insert<
+        IREE::Flow::FlowDialect, IREE::Util::UtilDialect, linalg::LinalgDialect,
+        mhlo::MhloDialect, shape::ShapeDialect, tensor::TensorDialect,
+        math::MathDialect, memref::MemRefDialect, complex::ComplexDialect>();
   }
 
   void runOnOperation() override {