[Mlir-commits] [mlir] [MLIR][TOSA] Guard scatter lowering against unranked operand (PR #178188)

[Ayush Kumar Gaur]

https://github.com/Ayush3941 created https://github.com/llvm/llvm-project/pull/178188

### What the problem 

--tosa-to-scf crashes when lowering tosa.scatter if operands are unranked, because the lowering assumes RankedTensorType and later builds tensor.extract_slice, which asserts on unranked tensors.

### Why it Happened

The current pattern does not check operand ranks before lowering, even though the implementation hardcodes 3D (N,W,C) semantics and cannot handle unranked or differently-ranked tensors.

### Whats The Fix

Add rank/type guards in ScatterOpConverter and fail the rewrite unless operands are ranked with expected ranks (values/input: rank-3, indices: rank-2), allowing legalization to fail gracefully instead of crashing.

>From 738d8cf3df00590d0371e785e6223d3bfdf5632d Mon Sep 17 00:00:00 2001
From: Ayush3941 <ayushkgaur1 at gmail.com>
Date: Tue, 27 Jan 2026 07:09:53 -0500
Subject: [PATCH] [MLIR][TOSA] Guard scatter lowering against unranked operand

---
 mlir/lib/Conversion/TosaToSCF/TosaToSCF.cpp             | 8 ++++++++
 mlir/test/Conversion/TosaToSCF/tosa-to-scf-invalid.mlir | 9 +++++++++
 2 files changed, 17 insertions(+)
 create mode 100644 mlir/test/Conversion/TosaToSCF/tosa-to-scf-invalid.mlir

diff --git a/mlir/lib/Conversion/TosaToSCF/TosaToSCF.cpp b/mlir/lib/Conversion/TosaToSCF/TosaToSCF.cpp
index aa6b4164e9876..9b1fb90d3ae17 100644
--- a/mlir/lib/Conversion/TosaToSCF/TosaToSCF.cpp
+++ b/mlir/lib/Conversion/TosaToSCF/TosaToSCF.cpp
@@ -101,6 +101,14 @@ class ScatterOpConverter : public OpRewritePattern<tosa::ScatterOp> {
     auto input = scatter.getInput();
     auto loc = scatter.getLoc();
 
+    auto valuesType  = dyn_cast<RankedTensorType>(valuesIn.getType());
+    auto indicesType = dyn_cast<RankedTensorType>(indices.getType());
+    auto inputType   = dyn_cast<RankedTensorType>(input.getType());
+    if (!valuesType || !indicesType || !inputType ||
+        valuesType.getRank() != 3 || inputType.getRank() != 3 ||
+        indicesType.getRank() != 2)
+      return failure();
+
     // N, W, C are chosen to match the TOSA spec
     auto dimN = createTensorDim(rewriter, loc, input, 0);
     auto dimW = createTensorDim(rewriter, loc, input, 1);
diff --git a/mlir/test/Conversion/TosaToSCF/tosa-to-scf-invalid.mlir b/mlir/test/Conversion/TosaToSCF/tosa-to-scf-invalid.mlir
new file mode 100644
index 0000000000000..c1bb458f42113
--- /dev/null
+++ b/mlir/test/Conversion/TosaToSCF/tosa-to-scf-invalid.mlir
@@ -0,0 +1,9 @@
+// RUN: mlir-opt --split-input-file --tosa-to-scf %s -verify-diagnostics -o -
+
+// CHECK-LABEL: @scatter_unranked
+func.func @scatter_unranked(%v: tensor<*xi32>, %idx: tensor<*xi32>, %inp: tensor<*xi32>) -> tensor<*xi32> {
+  // expected-error @+1 {{failed to legalize operation 'tosa.scatter' that was explicitly marked illegal}}
+  %0 = tosa.scatter %v, %idx, %inp
+      : (tensor<*xi32>, tensor<*xi32>, tensor<*xi32>) -> tensor<*xi32>
+  return %0 : tensor<*xi32>
+}