[Mlir-commits] [mlir] [mlir][tosa] Add missing int16 extension for row_gather (PR #206687)

[Luke Hutton]

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

Updates "availability" to include int16 extension to avoid assertion.

>From 27313a81c8b96de0f34707c0bdc4e26ee17dcb17 Mon Sep 17 00:00:00 2001
From: Luke Hutton <luke.hutton at arm.com>
Date: Tue, 30 Jun 2026 10:22:21 +0100
Subject: [PATCH] [mlir][tosa] Add missing int16 extension for row_gather

Updates "availability" to include int16 extension to avoid
assertion.

Change-Id: I6e660352ee35f67f1313f3f9d1399a5de2b0c2ee
---
 mlir/include/mlir/Dialect/Tosa/IR/TosaOps.td | 2 +-
 mlir/test/Dialect/Tosa/availability.mlir     | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/mlir/include/mlir/Dialect/Tosa/IR/TosaOps.td b/mlir/include/mlir/Dialect/Tosa/IR/TosaOps.td
index c01f32239a59f..8b065a4794c43 100644
--- a/mlir/include/mlir/Dialect/Tosa/IR/TosaOps.td
+++ b/mlir/include/mlir/Dialect/Tosa/IR/TosaOps.td
@@ -2596,7 +2596,7 @@ def Tosa_RowGatherOp : Tosa_InferShapedTypeOp<"row_gather", [NoMemoryEffect]> {
   list<Availability> availability = [Profile<[Tosa_PRO_INT, Tosa_PRO_FP]>,
                                      Extension<[Tosa_EXT_FP8E4M3,
                                                 Tosa_EXT_FP8E5M2, Tosa_EXT_BF16,
-                                                Tosa_EXT_INT64]>,
+                                                Tosa_EXT_INT16, Tosa_EXT_INT64]>,
   ];
 
   let hasCanonicalizer = 1;
diff --git a/mlir/test/Dialect/Tosa/availability.mlir b/mlir/test/Dialect/Tosa/availability.mlir
index d2a156e028fc8..f4f31d3f476e8 100644
--- a/mlir/test/Dialect/Tosa/availability.mlir
+++ b/mlir/test/Dialect/Tosa/availability.mlir
@@ -600,7 +600,7 @@ func.func @test_gather(%arg0: tensor<13x21x3xf32>, %arg1: tensor<13x26xi32>) ->
 func.func @test_row_gather(%arg0: tensor<13x21x3xf32>, %arg1: tensor<13x26xi32>) -> tensor<13x52x3xf32> {
   %row_count = "tosa.const"() {values = dense<2> : tensor<1xi32>} : () -> tensor<1xi32>
   // CHECK: profiles: [ [pro_int, pro_fp] ]
-  // CHECK: extensions: [ [fp8e4m3, fp8e5m2, bf16, int64] ]
+  // CHECK: extensions: [ [fp8e4m3, fp8e5m2, bf16, int16, int64] ]
   %0 = tosa.row_gather %arg0, %arg1, %row_count : (tensor<13x21x3xf32>, tensor<13x26xi32>, tensor<1xi32>) -> tensor<13x52x3xf32>
   return %0 : tensor<13x52x3xf32>
 }