[Mlir-commits] [mlir] [mlir][tosa] Fix validation of `dim` op when reliant on datatype extension (PR #180915)

[Luke Hutton]

https://github.com/lhutton1 created https://github.com/llvm/llvm-project/pull/180915

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.

>From 903711bcf995a0ad6de6e9274c58488aead90385 Mon Sep 17 00:00:00 2001
From: Luke Hutton <luke.hutton at arm.com>
Date: Wed, 11 Feb 2026 10:48:46 +0000
Subject: [PATCH] [mlir][tosa] Fix validation of `dim` op when reliant on
 datatype extension

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.

Change-Id: I5c986d6763904483c68da0e1c2aa9923947d2c3f
---
 mlir/include/mlir/Dialect/Tosa/IR/TosaShapeOps.td      |  5 +++++
 .../Tosa/tosa-validation-version-1p1-valid.mlir        | 10 ++++++++++
 2 files changed, 15 insertions(+)

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>