[Mlir-commits] [mlir] [mlir][tosa] Add mod_shape op (PR #170343)

[Luke Hutton]

https://github.com/lhutton1 updated https://github.com/llvm/llvm-project/pull/170343

>From 28f6c0cdb9aaef8077d098142928439749f538d6 Mon Sep 17 00:00:00 2001
From: Luke Hutton <luke.hutton at arm.com>
Date: Tue, 2 Dec 2025 17:44:41 +0000
Subject: [PATCH] [mlir][tosa] Add mod_shape op

Adds support for the mod_shape operation after
spec change: https://github.com/arm/tosa-specification/commit/efc88a100e2db06c2d6bc479fa63b26daab899ce.

This includes the operator definition, same rank
checks and level checks during validation. It does
not currently include support for folding or shape
inference. This will be added in a later commit.

Change-Id: I192332a8c7328fc41e7332b452bf540a1fb40d0f
---
 .../mlir/Dialect/Tosa/IR/TosaShapeOps.td       | 18 ++++++++++++++++++
 .../Tosa/Transforms/TosaProfileCompliance.cpp  |  1 +
 .../Dialect/Tosa/Transforms/TosaValidation.cpp |  1 +
 mlir/test/Dialect/Tosa/level_check.mlir        | 10 ++++++++++
 mlir/test/Dialect/Tosa/ops.mlir                |  9 +++++++++
 .../tosa-validation-version-1p1-valid.mlir     |  9 +++++++++
 6 files changed, 48 insertions(+)

diff --git a/mlir/include/mlir/Dialect/Tosa/IR/TosaShapeOps.td b/mlir/include/mlir/Dialect/Tosa/IR/TosaShapeOps.td
index d5a46b1b34312..389a44b252f5b 100644
--- a/mlir/include/mlir/Dialect/Tosa/IR/TosaShapeOps.td
+++ b/mlir/include/mlir/Dialect/Tosa/IR/TosaShapeOps.td
@@ -135,6 +135,24 @@ def Tosa_DivFloorShapeOp : Tosa_ElementwiseShapeOp<"div_floor_shape", [Pure]> {
   let results = (outs Tosa_Shape:$output);
 }
 
+//===----------------------------------------------------------------------===//
+// Operator: ModShape
+//===----------------------------------------------------------------------===//
+def Tosa_ModShapeOp : Tosa_ElementwiseShapeOp<"mod_shape", [Pure]> {
+  let summary = "Elementwise modulo of shapes.";
+
+  let description = [{
+    Elementwise modulo of input1 divided by input2.
+  }];
+
+  let arguments = (ins
+    Tosa_Shape:$input1,
+    Tosa_Shape:$input2
+  );
+
+  let results = (outs Tosa_Shape:$output);
+}
+
 //===----------------------------------------------------------------------===//
 // Operator: MulShape
 //===----------------------------------------------------------------------===//
diff --git a/mlir/lib/Dialect/Tosa/Transforms/TosaProfileCompliance.cpp b/mlir/lib/Dialect/Tosa/Transforms/TosaProfileCompliance.cpp
index c9150d5b34d00..06d1788075f8c 100644
--- a/mlir/lib/Dialect/Tosa/Transforms/TosaProfileCompliance.cpp
+++ b/mlir/lib/Dialect/Tosa/Transforms/TosaProfileCompliance.cpp
@@ -321,6 +321,7 @@ LogicalResult ProfileInfoDepot::populatationDispatch(Operation *op) {
   POPULATE_PROFILE_INFO_SKIP(ConstShape)
   POPULATE_PROFILE_INFO_SKIP(DivCeilShape)
   POPULATE_PROFILE_INFO_SKIP(DivFloorShape)
+  POPULATE_PROFILE_INFO_SKIP(ModShape)
   POPULATE_PROFILE_INFO_SKIP(MulShape)
   POPULATE_PROFILE_INFO_SKIP(SubShape)
   POPULATE_PROFILE_INFO_SKIP(Yield)
diff --git a/mlir/lib/Dialect/Tosa/Transforms/TosaValidation.cpp b/mlir/lib/Dialect/Tosa/Transforms/TosaValidation.cpp
index 530c6ae85287c..d2fe0f3f3a4f0 100644
--- a/mlir/lib/Dialect/Tosa/Transforms/TosaValidation.cpp
+++ b/mlir/lib/Dialect/Tosa/Transforms/TosaValidation.cpp
@@ -693,6 +693,7 @@ LogicalResult TosaValidation::levelCheckRanksAndSizes(Operation *op) {
   CHECK_RANKS(AddShape);
   CHECK_RANKS(DivCeilShape);
   CHECK_RANKS(DivFloorShape);
+  CHECK_RANKS(ModShape);
   CHECK_RANKS(MulShape);
   CHECK_RANKS(SubShape);
 
diff --git a/mlir/test/Dialect/Tosa/level_check.mlir b/mlir/test/Dialect/Tosa/level_check.mlir
index 213c4ae054c51..c337904f7b611 100644
--- a/mlir/test/Dialect/Tosa/level_check.mlir
+++ b/mlir/test/Dialect/Tosa/level_check.mlir
@@ -1682,3 +1682,13 @@ func.func @test_div_floor_shape_invalid_rank() -> !tosa.shape<7> {
   %c = tosa.div_floor_shape %a, %b : (!tosa.shape<7>, !tosa.shape<7>) -> !tosa.shape<7>
   return %c : !tosa.shape<7>
 }
+
+// -----
+
+func.func @test_mod_shape_invalid_rank() -> !tosa.shape<9> {
+  %a = tosa.const_shape {values = dense<[1, 2, 3, 4, 5, 6, 7, 8, 9]> : tensor<9xindex>} : () -> !tosa.shape<9>
+  %b = tosa.const_shape {values = dense<[1, 2, 3, 4, 5, 6, 7, 8, 9]> : tensor<9xindex>} : () -> !tosa.shape<9>
+  // expected-error at +1 {{'tosa.mod_shape' op failed shape type level check: '!tosa.shape<9>' exceeds MAX_RANK}}
+  %c = tosa.mod_shape %a, %b : (!tosa.shape<9>, !tosa.shape<9>) -> !tosa.shape<9>
+  return %c : !tosa.shape<9>
+}
diff --git a/mlir/test/Dialect/Tosa/ops.mlir b/mlir/test/Dialect/Tosa/ops.mlir
index b9e4d18156898..80f64761678cb 100644
--- a/mlir/test/Dialect/Tosa/ops.mlir
+++ b/mlir/test/Dialect/Tosa/ops.mlir
@@ -1409,6 +1409,15 @@ func.func @test_mul_shape() -> !tosa.shape<4> {
   return %c : !tosa.shape<4>
 }
 
+// -----
+// CHECK-LABEL: test_mod_shape
+func.func @test_mod_shape() -> !tosa.shape<4> {
+  %a = tosa.const_shape {values = dense<[9, 12, 10, 5]> : tensor<4xindex>} : () -> !tosa.shape<4>
+  %b = tosa.const_shape {values = dense<[2, 5, 3, 4]> : tensor<4xindex>} : () -> !tosa.shape<4>
+  %c = tosa.mod_shape %a, %b : (!tosa.shape<4>, !tosa.shape<4>) -> !tosa.shape<4>
+  return %c : !tosa.shape<4>
+}
+
 // -----
 // CHECK-LABEL: test_div_ceil_shape
 func.func @test_div_ceil_shape() -> !tosa.shape<4> {
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 10d322cf64fb7..474ba9635901b 100644
--- a/mlir/test/Dialect/Tosa/tosa-validation-version-1p1-valid.mlir
+++ b/mlir/test/Dialect/Tosa/tosa-validation-version-1p1-valid.mlir
@@ -158,3 +158,12 @@ func.func @test_add_shape() -> !tosa.shape<4> {
   %c = tosa.add_shape %a, %b : (!tosa.shape<4>, !tosa.shape<4>) -> !tosa.shape<4>
   return %c : !tosa.shape<4>
 }
+
+// -----
+// CHECK-LABEL: test_mod_shape
+func.func @test_mod_shape() -> !tosa.shape<3> {
+  %a = tosa.const_shape {values = dense<[10, 11, 12]> : tensor<3xindex>} : () -> !tosa.shape<3>
+  %b = tosa.const_shape {values = dense<[3, 5, 2]> : tensor<3xindex>} : () -> !tosa.shape<3>
+  %c = tosa.mod_shape %a, %b : (!tosa.shape<3>, !tosa.shape<3>) -> !tosa.shape<3>
+  return %c : !tosa.shape<3>
+}