[mlir][tosa] Fix validation of dim op when reliant on datatype extension

[lhutton1]

For example:

error: 'tosa.dim' op illegal: requires [bf16, shape] but not included in the profile compliance [shape]

    %0 = tosa.dim %arg0 {axis = 4 : i32} : (tensor<4x5x8x8x6x4xbf16>) -> !tosa.shape<1>

Here dim requires support to be declared for the BF16 and SHAPE extensions, but only SHAPE was specified in the op declaration.

[llvmbot]

@llvm/pr-subscribers-mlir-tosa

@llvm/pr-subscribers-mlir

Author: Luke Hutton (lhutton1)

Changes

For example:

error: 'tosa.dim' op illegal: requires [bf16, shape] but not included in the profile compliance [shape]

    %0 = tosa.dim %arg0 {axis = 4 : i32} : (tensor<4x5x8x8x6x4xbf16>) -> !tosa.shape<1>

Here dim requires support to be declared for the BF16 and SHAPE extensions, but only SHAPE was specified in the op declaration.

---

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

2 Files Affected:

• (modified) mlir/include/mlir/Dialect/Tosa/IR/TosaShapeOps.td (+5)
• (modified) mlir/test/Dialect/Tosa/tosa-validation-version-1p1-valid.mlir (+10)

diff --git a/mlir/include/mlir/Dialect/Tosa/IR/TosaShapeOps.td b/mlir/include/mlir/Dialect/Tosa/IR/TosaShapeOps.td
index cbcc2a017ac3a..25f7132fabea7 100644
--- a/mlir/include/mlir/Dialect/Tosa/IR/TosaShapeOps.td
+++ b/mlir/include/mlir/Dialect/Tosa/IR/TosaShapeOps.td
@@ -138,6 +138,11 @@ def Tosa_DimOp : Tosa_ShapeOp<"dim", [Pure]> {
 
   let results = (outs Tosa_Shape:$output);
 
+  list<Availability> availability = [
+    Profile<[Tosa_PRO_INT, Tosa_PRO_FP]>,
+    Extension<[Tosa_EXT_SHAPE, Tosa_EXT_FP8E4M3, Tosa_EXT_FP8E5M2, Tosa_EXT_BF16, Tosa_EXT_MXFP, Tosa_EXT_INT64]>,
+  ];
+
   let hasVerifier = 1;
 
   let hasFolder = 1;
diff --git a/mlir/test/Dialect/Tosa/tosa-validation-version-1p1-valid.mlir b/mlir/test/Dialect/Tosa/tosa-validation-version-1p1-valid.mlir
index 72269d21f3d98..fa03abcccacd8 100644
--- a/mlir/test/Dialect/Tosa/tosa-validation-version-1p1-valid.mlir
+++ b/mlir/test/Dialect/Tosa/tosa-validation-version-1p1-valid.mlir
@@ -271,6 +271,14 @@ func.func @test_dim(%arg0: tensor<1x2x3x4xi32>) -> !tosa.shape<1> {
   return %0 : !tosa.shape<1>
 }
 
+// -----
+
+// CHECK-LABEL: test_dim_bf16
+func.func @test_dim_bf16(%0: tensor<6x4x6x9xbf16>) {
+  %1 = tosa.dim %0 {axis = 1 : i32} : (tensor<6x4x6x9xbf16>) -> !tosa.shape<1>
+  return
+}
+
 // -----
 // CHECK-LABEL: test_exp2_shape
 func.func @test_exp2_shape() -> !tosa.shape<4> {
@@ -308,6 +316,8 @@ func.func @test_conv2d_block_scaled(%arg0: tensor<1x4x4x64xf4E2M1FN>, %arg1: ten
   return %0 : tensor<1x4x4x8xf32>
 }
 
+// -----
+
 // CHECK-LABEL: test_assert_equal_shape
 func.func @test_assert_equal_shape() {
   %0 = tosa.const_shape {values = dense<[10, 15]> : tensor<2xindex>} : () -> !tosa.shape<2>

[IanTaylerLessa-arm]

LGTM. Thank you for the quick fix!