[TorchConversion] Reject non-shaped function signatures in backend-co… (#4728)

The backend-contract verifiers check `func.func` legality with
`converter.isLegal(op)`, but that only looks at an op's operand/result
values. A function's argument and result types live in its
`FunctionType` attribute, which never gets inspected this way. So a
function argument of a non-shaped type that happens to be dead
(referenced by nothing in the body) sails through the verifier.

We hit this with an ONNX model whose graph input is an optional. The
importer preserves it as a `!torch.optional`argument,
`OptionalHasElement` constant-folds, and the argument is left dead on
the signature. The verifier reported success, and the illegal type then
blew up a downstream consumer instead of being caught here where it
should be. Tightening the verifier to also check the signature via
`isSignatureLegal` seems like the right place to stop this — it turns a
confusing downstream crash into the verifier's existing contract
diagnostic.

Same gap and fix apply to all three verifiers (linalg, TOSA, Stablehlo),
so I've fixed them together and added a regression test per backend.

One existing TOSA pipeline test had a dead `!torch.float` argument on
its signature and only passed because of this gap. TOSA operates on
tensors, not loose scalar primitives — had that scalar been used in the
body the verifier would already have rejected it. So the test was
locking down IR that isn't really legal; I dropped the unused argument.

---------

Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>

Author: hariprasadravi

lib/Dialect/TorchConversion/Transforms/VerifyLinalgOnTensorsBackendContract.cpp

@@ -56,6 +56,12 @@ class VerifyLinalgOnTensorsBackendContractPass
     }
 
     auto opHasLegalTypes = [&](Operation *op) { return converter.isLegal(op); };
+    // `converter.isLegal(op)` inspects only operands/results; a `func.func`'s
+    // signature is in its `FunctionType` attribute and would otherwise go
+    // unchecked (e.g. an unused non-shaped arg like `!torch.optional`).
+    auto funcHasLegalTypes = [&](func::FuncOp func) {
+      return converter.isSignatureLegal(func.getFunctionType());
+    };
     auto isLegalScalarOp = [&](Operation *op) {
       // We recognize basic scalar ops by them having the trait "Elementwise",
       // even though we don't expect them to operate on tensors.
@@ -66,8 +72,9 @@ class VerifyLinalgOnTensorsBackendContractPass
     ConversionTarget target(*context);
 
     // Structural operations.
-    target.addDynamicallyLegalOp<ModuleOp, func::FuncOp, func::ReturnOp>(
-        opHasLegalTypes);
+    target.addDynamicallyLegalOp<func::FuncOp>(funcHasLegalTypes);
+    target.addDynamicallyLegalOp<func::ReturnOp>(opHasLegalTypes);
+    target.addLegalOp<ModuleOp>();
 
     target.addDynamicallyLegalOp<GetNextSeedOp>(opHasLegalTypes);
 

lib/Dialect/TorchConversion/Transforms/VerifyStablehloBackendContract.cpp

@@ -42,13 +42,20 @@ class VerifyStablehloBackendContractPass
     });
 
     auto opHasLegalTypes = [&](Operation *op) { return converter.isLegal(op); };
+    // `converter.isLegal(op)` inspects only operands/results; a `func.func`'s
+    // signature is in its `FunctionType` attribute and would otherwise go
+    // unchecked (e.g. an unused non-shaped arg like `!torch.optional`).
+    auto funcHasLegalTypes = [&](func::FuncOp func) {
+      return converter.isSignatureLegal(func.getFunctionType());
+    };
 
     MLIRContext *context = &getContext();
     ConversionTarget target(*context);
 
     // Structural operations.
-    target.addDynamicallyLegalOp<ModuleOp, func::FuncOp, func::ReturnOp>(
-        opHasLegalTypes);
+    target.addDynamicallyLegalOp<func::FuncOp>(funcHasLegalTypes);
+    target.addDynamicallyLegalOp<func::ReturnOp>(opHasLegalTypes);
+    target.addLegalOp<ModuleOp>();
 
     target.addLegalDialect<chlo::ChloDialect>();
     target.addLegalDialect<stablehlo::StablehloDialect>();

lib/Dialect/TorchConversion/Transforms/VerifyTosaBackendContract.cpp

@@ -38,12 +38,19 @@ class VerifyTosaBackendContractPass
     });
 
     auto opHasLegalTypes = [&](Operation *op) { return converter.isLegal(op); };
+    // `converter.isLegal(op)` inspects only operands/results; a `func.func`'s
+    // signature is in its `FunctionType` attribute and would otherwise go
+    // unchecked (e.g. an unused non-shaped arg like `!torch.optional`).
+    auto funcHasLegalTypes = [&](func::FuncOp func) {
+      return converter.isSignatureLegal(func.getFunctionType());
+    };
 
     ConversionTarget target(*context);
 
     // Structural operations.
-    target.addDynamicallyLegalOp<ModuleOp, func::FuncOp, func::ReturnOp>(
-        opHasLegalTypes);
+    target.addDynamicallyLegalOp<func::FuncOp>(funcHasLegalTypes);
+    target.addDynamicallyLegalOp<func::ReturnOp>(opHasLegalTypes);
+    target.addLegalOp<ModuleOp>();
     // Basic scalar operations.
     target.addLegalDialect<tosa::TosaDialect>();
     target.addDynamicallyLegalOp<tensor::CastOp>(opHasLegalTypes);

test/Conversion/TorchToTosa/torch-backend-to-tosa-backend-pipeline.mlir

@@ -20,7 +20,7 @@ func.func @torch.aten.mul.Scalar$mixed_type(%arg0: !torch.vtensor<[5],bf16>) ->
 // CHECK-SAME: %[[VAL_1:.*]]: tensor<6xf32>
 // CHECK: %[[VAL_3:.*]] = tosa.cast %[[VAL_1]] : (tensor<6xf32>) -> tensor<6xbf16>
 // CHECK: %[[VAL_4:.*]] = tosa.add %[[VAL_0]], %[[VAL_3]] : (tensor<6xbf16>, tensor<6xbf16>) -> tensor<6xbf16>
-func.func @torch.aten.add.Tensor$mixed_type_fp(%arg0: !torch.vtensor<[6],bf16>, %arg1: !torch.vtensor<[6],f32>, %arg2: !torch.float) -> !torch.vtensor<[6],bf16> {
+func.func @torch.aten.add.Tensor$mixed_type_fp(%arg0: !torch.vtensor<[6],bf16>, %arg1: !torch.vtensor<[6],f32>) -> !torch.vtensor<[6],bf16> {
   %float1 = torch.constant.float 1.000000e+00
   %0 = torch.aten.add.Tensor %arg0, %arg1, %float1 : !torch.vtensor<[6],bf16>, !torch.vtensor<[6],f32>, !torch.float -> !torch.vtensor<[6],bf16>
   return %0 : !torch.vtensor<[6],bf16>

test/Dialect/TorchConversion/verify-linalg-on-tensors-backend-contract.mlir

@@ -46,8 +46,22 @@ module {
 
 // expected-error@+1 {{Module does not conform to the linalg-on-tensors backend contract.}}
 module {
+  // expected-error@+1 {{failed to legalize operation 'func.func'}}
   func.func @disallowed(%arg0: !torch.tensor) -> !torch.tensor {
-    // expected-error@+1 {{failed to legalize operation 'func.return'}}
     return %arg0 : !torch.tensor
   }
 }
+
+// -----
+
+// A non-shaped function argument (e.g. `!torch.optional`) that is dead --
+// referenced by no op in the body -- should be rejected.
+
+// expected-error@+1 {{Module does not conform to the linalg-on-tensors backend contract.}}
+module {
+  // expected-error@+1 {{failed to legalize operation 'func.func'}}
+  func.func @dead_optional_arg(%arg0: !torch.optional<vtensor<[4],f32>>) -> tensor<1xi64> {
+    %cst = arith.constant dense<1> : tensor<1xi64>
+    return %cst : tensor<1xi64>
+  }
+}

test/Dialect/TorchConversion/verify-stablehlo-backend-contract.mlir

@@ -0,0 +1,35 @@
+// RUN: torch-mlir-opt -torch-verify-stablehlo-backend-contract -split-input-file -verify-diagnostics -allow-unregistered-dialect %s | FileCheck %s
+// REQUIRES: stablehlo
+
+// CHECK: func.func @tanh
+func.func @tanh(%arg0: tensor<?x?xf32>) -> tensor<?x?xf32> {
+  %0 = stablehlo.tanh %arg0 : tensor<?x?xf32>
+  return %0 : tensor<?x?xf32>
+}
+
+// -----
+
+// Basic check of error reporting.
+
+// expected-error@+1 {{Module does not conform to the Stablehlo backend contract.}}
+module {
+  func.func @disallowed() {
+    // expected-error@+1 {{failed to legalize operation 'unknown_dialect.unknown_op'}}
+    "unknown_dialect.unknown_op"() : () -> ()
+    return
+  }
+}
+
+// -----
+
+// A non-shaped function argument (e.g. `!torch.optional`) that is dead --
+// referenced by no op in the body -- should be rejected.
+
+// expected-error@+1 {{Module does not conform to the Stablehlo backend contract.}}
+module {
+  // expected-error@+1 {{failed to legalize operation 'func.func'}}
+  func.func @dead_optional_arg(%arg0: !torch.optional<vtensor<[4],f32>>) -> tensor<1xi64> {
+    %cst = arith.constant dense<1> : tensor<1xi64>
+    return %cst : tensor<1xi64>
+  }
+}

test/Dialect/TorchConversion/verify-tosa-backend-contract.mlir

@@ -35,8 +35,22 @@ module {
 
 // expected-error@+1 {{Module does not conform to the TOSA backend contract.}}
 module {
+  // expected-error@+1 {{failed to legalize operation 'func.func'}}
   func.func @disallowed(%arg0: !torch.tensor) -> !torch.tensor {
-    // expected-error@+1 {{failed to legalize operation 'func.return'}}
     return %arg0 : !torch.tensor
   }
 }
+
+// -----
+
+// A non-shaped function argument (e.g. `!torch.optional`) that is dead --
+// referenced by no op in the body -- should be rejected.
+
+// expected-error@+1 {{Module does not conform to the TOSA backend contract.}}
+module {
+  // expected-error@+1 {{failed to legalize operation 'func.func'}}
+  func.func @dead_optional_arg(%arg0: !torch.optional<vtensor<[4],f32>>) -> tensor<1xi64> {
+    %cst = arith.constant dense<1> : tensor<1xi64>
+    return %cst : tensor<1xi64>
+  }
+}

test/lit.cfg.py

@@ -35,6 +35,11 @@
 config.substitutions.append(("%PATH%", config.environment["PATH"]))
 config.substitutions.append(("%shlibext", config.llvm_shlib_ext))
 
+# Register optional-backend availability as lit features so tests that exercise a
+# specific backend contract can guard themselves with `// REQUIRES: <backend>`.
+if getattr(config, "enable_stablehlo", False):
+    config.available_features.add("stablehlo")
+
 llvm_config.with_system_environment(["HOME", "INCLUDE", "LIB", "TMP", "TEMP"])
 
 # llvm_config.use_default_substitutions()