[Mlir-commits] [mlir] [mlir][tosa] Preserve non-TOSA ops in narrowing passes (PR #224594)

[Jack Vickers]

https://github.com/jackv1ck updated https://github.com/llvm/llvm-project/pull/224594

>From 2de2087fb786d2714f437130d2d4f72cfccc98e8 Mon Sep 17 00:00:00 2001
From: Jack Vickers <jack.vickers at arm.com>
Date: Tue, 15 Sep 2026 16:17:12 +0000
Subject: [PATCH] [mlir][tosa] Preserve non-TOSA ops in narrowing passes

The TOSA narrowing passes currently reject operations from non-TOSA dialects even when their types do not require conversion.

Mark unknown ops as legal provided that their operand and result types are legal.
This follows the approach used by TosaConvertIntegerTypeToSignless and allows mixed-dialect IR to pass through TOSA narrowing.

Update the rollback lit test to continue forcing a conversion failure using illegal operand and result types.

Assisted-by: Codex
---
 .../Tosa/Transforms/TosaNarrowTypes.cpp       | 26 +++++++++++--------
 .../Dialect/Tosa/tosa-narrow-f32-to-f16.mlir  |  9 +++++++
 .../Dialect/Tosa/tosa-narrow-f64-to-f32.mlir  |  9 +++++++
 .../Dialect/Tosa/tosa-narrow-i64-to-i32.mlir  | 15 ++++++++---
 4 files changed, 45 insertions(+), 14 deletions(-)

diff --git a/mlir/lib/Dialect/Tosa/Transforms/TosaNarrowTypes.cpp b/mlir/lib/Dialect/Tosa/Transforms/TosaNarrowTypes.cpp
index c53704965ed09..70fb94e386b09 100644
--- a/mlir/lib/Dialect/Tosa/Transforms/TosaNarrowTypes.cpp
+++ b/mlir/lib/Dialect/Tosa/Transforms/TosaNarrowTypes.cpp
@@ -751,18 +751,21 @@ LogicalResult runTosaNarrowing(Operation *op, bool aggressiveRewrite,
         });
   }
 
+  auto hasLegalTypes = [&typeConverter](Operation *op) {
+    return typeConverter.isLegal(op->getOperandTypes()) &&
+           typeConverter.isLegal(op->getResultTypes());
+  };
+
   ConversionTarget target(*context);
-  target.addDynamicallyLegalDialect<tosa::TosaDialect>(
-      [&typeConverter, convertAccumulatorType](Operation *op) {
-        if (!typeConverter.isLegal(op->getResultTypes()) ||
-            !typeConverter.isLegal(op->getOperandTypes()))
-          return false;
-        if (!convertAccumulatorType)
-          return true;
-        const auto accumulatorType = op->getAttrOfType<TypeAttr>("acc_type");
-        return !accumulatorType ||
-               !typeNeedsConversion<Kind>(accumulatorType.getValue());
-      });
+  target.addDynamicallyLegalDialect<tosa::TosaDialect>([&](Operation *op) {
+    if (!hasLegalTypes(op))
+      return false;
+    if (!convertAccumulatorType)
+      return true;
+    const auto accumulatorType = op->getAttrOfType<TypeAttr>("acc_type");
+    return !accumulatorType ||
+           !typeNeedsConversion<Kind>(accumulatorType.getValue());
+  });
   if (convertFunctionBoundaries) {
     target.addDynamicallyLegalOp<func::FuncOp>(
         [&typeConverter](func::FuncOp op) {
@@ -780,6 +783,7 @@ LogicalResult runTosaNarrowing(Operation *op, bool aggressiveRewrite,
     target.addDynamicallyLegalOp<func::ReturnOp>(
         [](func::ReturnOp) { return true; });
   }
+  target.markUnknownOpDynamicallyLegal(hasLegalTypes);
 
   RewritePatternSet patterns(context);
   if (convertFunctionBoundaries) {
diff --git a/mlir/test/Dialect/Tosa/tosa-narrow-f32-to-f16.mlir b/mlir/test/Dialect/Tosa/tosa-narrow-f32-to-f16.mlir
index 3c3bd2121de14..5f55ecbcb84ab 100644
--- a/mlir/test/Dialect/Tosa/tosa-narrow-f32-to-f16.mlir
+++ b/mlir/test/Dialect/Tosa/tosa-narrow-f32-to-f16.mlir
@@ -221,3 +221,12 @@ func.func @test_dense_resource_f32() -> tensor<1x2xf32> {
     }
   }
 #-}
+
+// -----
+
+// CHECK-LABEL: test_preserve_unknown_dialect_op
+func.func @test_preserve_unknown_dialect_op() -> i32 {
+  //COMMON: arith.constant 1 : i32
+  %0 = arith.constant 1 : i32
+  return %0 : i32
+}
diff --git a/mlir/test/Dialect/Tosa/tosa-narrow-f64-to-f32.mlir b/mlir/test/Dialect/Tosa/tosa-narrow-f64-to-f32.mlir
index 7c78f0067ffca..ccdb340c090ff 100644
--- a/mlir/test/Dialect/Tosa/tosa-narrow-f64-to-f32.mlir
+++ b/mlir/test/Dialect/Tosa/tosa-narrow-f64-to-f32.mlir
@@ -198,3 +198,12 @@ func.func @test_dense_ressource_f64() -> tensor<1x2xf64> {
     }
   }
 #-}
+
+// -----
+
+// CHECK-LABEL: test_preserve_unknown_dialect_op
+func.func @test_preserve_unknown_dialect_op() -> i32 {
+  //COMMON: arith.constant 1 : i32
+  %0 = arith.constant 1 : i32
+  return %0 : i32
+}
\ No newline at end of file
diff --git a/mlir/test/Dialect/Tosa/tosa-narrow-i64-to-i32.mlir b/mlir/test/Dialect/Tosa/tosa-narrow-i64-to-i32.mlir
index a8b5de834de8c..4fc5677591c16 100644
--- a/mlir/test/Dialect/Tosa/tosa-narrow-i64-to-i32.mlir
+++ b/mlir/test/Dialect/Tosa/tosa-narrow-i64-to-i32.mlir
@@ -77,8 +77,8 @@ func.func @test_regions(%arg0: tensor<1x2xi32>, %arg1: tensor<1xi32>, %arg2: ten
 
 // CHECK-LABEL: test_cond_if_i64_yield_with_call
 module {
-  func.func @m0() -> () {
-    return
+  func.func @m0(%arg0: tensor<4xi64>) -> tensor<4xi64> {
+    return %arg0 : tensor<4xi64>
   }
   func.func @test_cond_if_i64_yield_with_call(%arg0: tensor<1xi1>, %arg1: tensor<4xi64>, %arg2: tensor<4xi64>) -> () {
     %0 = tosa.cond_if %arg0 (%arg3 = %arg1) : tensor<1xi1> (tensor<4xi64>) -> tensor<4xi64> {
@@ -89,7 +89,7 @@ module {
       tosa.yield %arg3 : tensor<4xi64>
     }
     // expected-error @+1 {{failed to legalize operation 'func.call'}}
-    call @m0() : () -> ()
+    %1 = call @m0(%arg1) : (tensor<4xi64>) -> tensor<4xi64>
     return
   }
 }
@@ -254,3 +254,12 @@ func.func @test_dense_ressource_i64() -> tensor<1x2xi64> {
     }
   }
 #-}
+
+// -----
+
+// CHECK-LABEL: test_preserve_unknown_dialect_op
+func.func @test_preserve_unknown_dialect_op() -> i32 {
+  // COMMON: arith.constant 1 : i32
+  %0 = arith.constant 1 : i32
+  return %0 : i32
+}