Triggering auto torch input conversion based on func arg/result types. (#21067)

Fixes #21066.

Author: benvanik

compiler/plugins/input/Torch/InputConversion/test/auto_input_conversion.mlir

@@ -2,10 +2,20 @@
 
 // Check that the auto input conversion pipeline uses this plugin.
 
-
 // CHECK-LABEL: util.func public @simple_add_onnx
 // CHECK:  arith.addi
 func.func @simple_add_onnx(%arg0: !torch.vtensor<[],si64>, %arg1: !torch.vtensor<[],si64>) -> !torch.vtensor<[],si64> attributes {torch.onnx_meta.ir_version = 8 : si64, torch.onnx_meta.opset_version = 17 : si64, torch.onnx_meta.producer_name = "pytorch", torch.onnx_meta.producer_version = "2.1.0"} {
   %0 = torch.operator "onnx.Add"(%arg0, %arg1) : (!torch.vtensor<[],si64>, !torch.vtensor<[],si64>) -> !torch.vtensor<[],si64>
   return %0 : !torch.vtensor<[],si64>
 }
+
+// -----
+
+// Tests that a function using torch types but not containing any ops is still
+// handled by the torch input pipeline.
+
+// CHECK: util.func public @nop$async
+// CHECK: util.func public @nop(%{{.+}}: !hal.buffer_view) -> !hal.buffer_view
+func.func @nop(%arg0: !torch.vtensor<[5],f32>) -> !torch.vtensor<[5],f32> attributes {torch.assume_strict_symbolic_shapes} {
+  return %arg0 : !torch.vtensor<[5],f32>
+}

compiler/plugins/input/Torch/PluginRegistration.cpp

@@ -129,11 +129,18 @@ struct TorchSession
       return WalkResult::advance();
     });
 
+    auto isTorchType = [&](Type type) {
+      return &type.getDialect() == torchDialect;
+    };
     for (auto funcOp : module.getOps<func::FuncOp>()) {
+      auto funcType = funcOp.getFunctionType();
+      if (llvm::any_of(funcType.getInputs(), isTorchType) ||
+          llvm::any_of(funcType.getResults(), isTorchType)) {
+        hasTorch = true;
+      }
       if (funcOp->getAttrOfType<mlir::IntegerAttr>(
               "torch.onnx_meta.opset_version")) {
         hasOnnx = true;
-        break;
       }
     }