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

[llvmlistbot at llvm.org]

Author: Jack Vickers
Date: 2026-09-29T12:57:23+01:00
New Revision: 39b72080e0399a15b0313b4eb24e3dbb0aee3c08

URL: https://github.com/llvm/llvm-project/commit/39b72080e0399a15b0313b4eb24e3dbb0aee3c08
DIFF: https://github.com/llvm/llvm-project/commit/39b72080e0399a15b0313b4eb24e3dbb0aee3c08.diff

LOG: [mlir][tosa] Preserve non-TOSA ops in narrowing passes (#224594)

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

This change marks 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.

Adds a lit test for each narrowing pass to prove that mixed-dialect is
allowed.

Updates the i64-to-i32 rollback lit test to continue forcing a
conversion failure using illegal operand and result types.

Assisted by: Codex

Added: 
    

Modified: 
    mlir/lib/Dialect/Tosa/Transforms/TosaNarrowTypes.cpp
    mlir/test/Dialect/Tosa/tosa-narrow-f32-to-f16.mlir
    mlir/test/Dialect/Tosa/tosa-narrow-f64-to-f32.mlir
    mlir/test/Dialect/Tosa/tosa-narrow-i64-to-i32.mlir

Removed: 
    


################################################################################
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 13ec6576dcb92..e822f67518d12 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 a5fdcf64da477..db1a62b201e4f 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 415728fcb8792..2c787eec4f25c 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
+}