[TOSA] Make validation pass isValidElementType check more strict

[lhutton1]

The validation pass is used to check alignment of the IR against the TOSA specification. This commit updates the isValidElement check to more strictly align with the specifications supported element types.

[llvmbot]

@llvm/pr-subscribers-mlir

@llvm/pr-subscribers-mlir-tosa

Author: Luke Hutton (lhutton1)

Changes

The validation pass is used to check alignment of the IR against the TOSA specification. This commit updates the isValidElement check to more strictly align with the specifications supported element types.

---

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

2 Files Affected:

• (modified) mlir/lib/Dialect/Tosa/Transforms/TosaValidation.cpp (+3-16)
• (modified) mlir/test/Dialect/Tosa/level_check.mlir (+8)

diff --git a/mlir/lib/Dialect/Tosa/Transforms/TosaValidation.cpp b/mlir/lib/Dialect/Tosa/Transforms/TosaValidation.cpp
index 93e8cac6b84e9c..893cedefc1ebde 100644
--- a/mlir/lib/Dialect/Tosa/Transforms/TosaValidation.cpp
+++ b/mlir/lib/Dialect/Tosa/Transforms/TosaValidation.cpp
@@ -524,18 +524,8 @@ bool TosaValidation::isValidElementType(Type type) {
     if (!isEnabledProfile(TosaProfileEnum::MainInference))
       return false;
     return type.isF32() || type.isF16() || type.isBF16();
-  }
-  if (auto intTy = dyn_cast<IntegerType>(type)) {
-    if (intTy.isUnsigned()) {
-      switch (intTy.getWidth()) {
-      case 8:
-      case 16:
-        return true;
-      default:
-        return false;
-      }
-    } else {
-      // Signless - treated as signed.
+  } else if (auto intTy = dyn_cast<IntegerType>(type)) {
+    if (intTy.isSignless()) {
       switch (intTy.getWidth()) {
       case 1:
       case 4:
@@ -544,13 +534,10 @@ bool TosaValidation::isValidElementType(Type type) {
       case 32:
       case 48:
         return true;
-      default:
-        return false;
       }
     }
-    return false;
   }
-  return true;
+  return false;
 }
 
 void TosaValidation::runOnOperation() {
diff --git a/mlir/test/Dialect/Tosa/level_check.mlir b/mlir/test/Dialect/Tosa/level_check.mlir
index e851019362958f..529a16ca48c7eb 100644
--- a/mlir/test/Dialect/Tosa/level_check.mlir
+++ b/mlir/test/Dialect/Tosa/level_check.mlir
@@ -143,6 +143,14 @@ func.func @test_const_f64(%arg0 : tensor<1xf64>) {
 
 // -----
 
+func.func @test_const_ui8(%arg0 : tensor<1xui8>) {
+  // expected-error@+1 {{'tosa.const' op is not profile-aligned: element type 'ui8' is not legal}}
+  %0 = "tosa.const"() {value = dense<0> : tensor<1xui8>} : () -> tensor<1xui8>
+  return
+}
+
+// -----
+
 func.func @test_avgpool2d_kernel_y(%arg0: tensor<1x32x32x8xf32>) -> tensor<1x32x32x8xf32> {
   // expected-error@+1 {{'tosa.avg_pool2d' op failed level check: kernel <= MAX_KERNEL}}
   %0 = "tosa.avg_pool2d"(%arg0) {kernel = array<i64: 8193, 1>, pad = array<i64: 4, 4, 4, 4>, stride = array<i64: 1, 1>, acc_type = f32} :

[GeorgeARM]

LGTM.
I recall that rescale could return u8 back but this is already fixed from what I see.

[MaheshRavishankar]

This seems to fail on IndexType. Is that explicit.

(Found as an error while integrate it into downstream iree-org/iree#19479 )

[lhutton1]

Thanks for reporting and for the revert @MaheshRavishankar, apologies I'm only just taking a look now.

I took a look at the failures and noticed that the validation pass is being run over mixed dialect IR. Currently the pass seems to run over all operators regardless of the dialect, including operators from arith/tensor/linalg that use the index type which would be considered illegal in tosa. I suspect we'll need to ensure only tosa operators are checked by this pass - I'll work on a patch for this, then resubmit this change

[lhutton1]

I thought this commit had been reverted in LLVM, but it seems I was mistaken and it was reverted in the IREE fork instead. I'm unsure of the official process here, but I wanted to give a quick update. Since #120205 was merged, my hope is that the failures in IREE have been fixed - please let me know if that's not the case so I can investigate further

[MaheshRavishankar]

I thought this commit had been reverted in LLVM, but it seems I was mistaken and it was reverted in the IREE fork instead. I'm unsure of the official process here, but I wanted to give a quick update. Since #120205 was merged, my hope is that the failures in IREE have been fixed - please let me know if that's not the case so I can investigate further

Thanks @lhutton1 . @raikonenfnu if you get around to another integrate of LLVM into IREE could you please try dropping the revert in IREEs LLVM fork