[mlir][TosaToLinalg] Fix TosaToLinalg to restrict tosa.cast types to integer or float

[CoTinker]

This PR fixes a bug where TosaToLinalg incorrectly allows tosa.cast to accept types other than integer or float.
Fixes #116342.

[llvmbot]

@llvm/pr-subscribers-mlir-tosa

@llvm/pr-subscribers-mlir

Author: Longsheng Mou (CoTinker)

Changes

This PR fixes a bug where TosaToLinalg incorrectly allows tosa.cast to accept types other than integer or float.
Fixes #116342.

---

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

2 Files Affected:

• (modified) mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp (+5)
• (modified) mlir/test/Conversion/TosaToLinalg/tosa-to-linalg-invalid.mlir (+8)

diff --git a/mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp b/mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp
index 607667fcc6945..e5994cdc777b2 100644
--- a/mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp
+++ b/mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp
@@ -524,6 +524,11 @@ static Value createLinalgBodyCalculationForElementwiseOp(
   if (isa<tosa::CastOp>(op)) {
     Type srcTy = elementTy;
     Type dstTy = resultTypes.front();
+    if (!srcTy.isIntOrFloat() || !dstTy.isIntOrFloat()) {
+      (void)rewriter.notifyMatchFailure(op,"unsupported type");
+      return nullptr;
+    }
+
     bool bitExtend =
         srcTy.getIntOrFloatBitWidth() < dstTy.getIntOrFloatBitWidth();
 
diff --git a/mlir/test/Conversion/TosaToLinalg/tosa-to-linalg-invalid.mlir b/mlir/test/Conversion/TosaToLinalg/tosa-to-linalg-invalid.mlir
index 460e207d62de6..5db3f56cf459e 100644
--- a/mlir/test/Conversion/TosaToLinalg/tosa-to-linalg-invalid.mlir
+++ b/mlir/test/Conversion/TosaToLinalg/tosa-to-linalg-invalid.mlir
@@ -54,3 +54,11 @@ func.func @test_add_2d_different_ranks(%arg0: tensor<3x4xf32>, %arg1: tensor<2x3
   %0 = "tosa.add"(%arg0, %arg1) : (tensor<3x4xf32>, tensor<2x3x4xf32>) -> tensor<2x3x4xf32>
   return %0 : tensor<2x3x4xf32>
 }
+
+// -----
+
+func.func @cast_unsupported_type(%arg0: tensor<13x21x3xi32>) -> tensor<13x21x3x!quant.uniform<i16:f32, 0.078431375324726104:128>> {
+  // expected-error@+1 {{failed to legalize operation 'tosa.cast'}}
+  %0 = tosa.cast %arg0 : (tensor<13x21x3xi32>) -> tensor<13x21x3x!quant.uniform<i16:f32, 0.078431375324726104:128>>
+  return %0 : tensor<13x21x3x!quant.uniform<i16:f32, 0.078431375324726104:128>>
+}

[llvmbot]

@llvm/pr-subscribers-mlir-linalg

Author: Longsheng Mou (CoTinker)

Changes

This PR fixes a bug where TosaToLinalg incorrectly allows tosa.cast to accept types other than integer or float.
Fixes #116342.

---

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

2 Files Affected:

• (modified) mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp (+5)
• (modified) mlir/test/Conversion/TosaToLinalg/tosa-to-linalg-invalid.mlir (+8)

diff --git a/mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp b/mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp
index 607667fcc6945..e5994cdc777b2 100644
--- a/mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp
+++ b/mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp
@@ -524,6 +524,11 @@ static Value createLinalgBodyCalculationForElementwiseOp(
   if (isa<tosa::CastOp>(op)) {
     Type srcTy = elementTy;
     Type dstTy = resultTypes.front();
+    if (!srcTy.isIntOrFloat() || !dstTy.isIntOrFloat()) {
+      (void)rewriter.notifyMatchFailure(op,"unsupported type");
+      return nullptr;
+    }
+
     bool bitExtend =
         srcTy.getIntOrFloatBitWidth() < dstTy.getIntOrFloatBitWidth();
 
diff --git a/mlir/test/Conversion/TosaToLinalg/tosa-to-linalg-invalid.mlir b/mlir/test/Conversion/TosaToLinalg/tosa-to-linalg-invalid.mlir
index 460e207d62de6..5db3f56cf459e 100644
--- a/mlir/test/Conversion/TosaToLinalg/tosa-to-linalg-invalid.mlir
+++ b/mlir/test/Conversion/TosaToLinalg/tosa-to-linalg-invalid.mlir
@@ -54,3 +54,11 @@ func.func @test_add_2d_different_ranks(%arg0: tensor<3x4xf32>, %arg1: tensor<2x3
   %0 = "tosa.add"(%arg0, %arg1) : (tensor<3x4xf32>, tensor<2x3x4xf32>) -> tensor<2x3x4xf32>
   return %0 : tensor<2x3x4xf32>
 }
+
+// -----
+
+func.func @cast_unsupported_type(%arg0: tensor<13x21x3xi32>) -> tensor<13x21x3x!quant.uniform<i16:f32, 0.078431375324726104:128>> {
+  // expected-error@+1 {{failed to legalize operation 'tosa.cast'}}
+  %0 = tosa.cast %arg0 : (tensor<13x21x3xi32>) -> tensor<13x21x3x!quant.uniform<i16:f32, 0.078431375324726104:128>>
+  return %0 : tensor<13x21x3x!quant.uniform<i16:f32, 0.078431375324726104:128>>
+}

[github-actions]

✅ With the latest revision this PR passed the C/C++ code formatter.

[Jerry-Ge]

LGTM

[Jerry-Ge]

Just saw comments here regarding this issue (https://discourse.llvm.org/t/tosa-cast-valid-supported-data-types/84808). I think we should abandon this PR now. What do you think? @CoTinker

[GeorgeARM]

This relates to the TOSA legalizations and casting of types like MLIR quantized ones that linalg will probably fail/misbehave.
So the patch looks fine from what I gather? Unless I am missing something @Jerry-Ge

[Jerry-Ge]

This relates to the TOSA legalizations and casting of types like MLIR quantized ones that linalg will probably fail/misbehave. So the patch looks fine from what I gather? Unless I am missing something @Jerry-Ge

Quote from @sjarus 's comment here: https://discourse.llvm.org/t/tosa-cast-valid-supported-data-types/84808

Invoking this pass between TOSA generation and downstream compilation should capture cases where spec compliance requires changes to the legalization. Broadly disabling any spec non-compliant legalizations might not be a good approach, since it blocks experimental activity.

[GeorgeARM]

Sure but the legalization cannot handle it properly so is failing. If the legalization supports it then yes we can enable and leave the validation pass for compliance checking.

[Jerry-Ge]

Sure but the legalization cannot handle it properly so is failing. If the legalization supports it then yes we can enable and leave the validation pass for compliance checking.

I agree. Let's merge it.