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

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir-tosa

Author: Jack Vickers (jackv1ck)

<details>
<summary>Changes</summary>

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.



---
Full diff: https://github.com/llvm/llvm-project/pull/224594.diff


4 Files Affected:

- (modified) mlir/lib/Dialect/Tosa/Transforms/TosaNarrowTypes.cpp (+15-11) 
- (modified) mlir/test/Dialect/Tosa/tosa-narrow-f32-to-f16.mlir (+9) 
- (modified) mlir/test/Dialect/Tosa/tosa-narrow-f64-to-f32.mlir (+9) 
- (modified) mlir/test/Dialect/Tosa/tosa-narrow-i64-to-i32.mlir (+12-3) 


``````````diff
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
+}

``````````

</details>


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