[Mlir-commits] [mlir] [mlir][tosa] Add mod_shape op (PR #170343)
Luke Hutton
llvmlistbot at llvm.org
Thu Jan 8 13:25:13 PST 2026
https://github.com/lhutton1 updated https://github.com/llvm/llvm-project/pull/170343
>From bc6103eb57fcc6d33cdab1b3661847002ee997fc 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 1/2] [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 | 10 ++++++++++
6 files changed, 49 insertions(+)
diff --git a/mlir/include/mlir/Dialect/Tosa/IR/TosaShapeOps.td b/mlir/include/mlir/Dialect/Tosa/IR/TosaShapeOps.td
index 79967b7c9585e..f4ac59ebdbdba 100644
--- a/mlir/include/mlir/Dialect/Tosa/IR/TosaShapeOps.td
+++ b/mlir/include/mlir/Dialect/Tosa/IR/TosaShapeOps.td
@@ -175,6 +175,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 e8a24057a96ac..0b7925549771f 100644
--- a/mlir/lib/Dialect/Tosa/Transforms/TosaProfileCompliance.cpp
+++ b/mlir/lib/Dialect/Tosa/Transforms/TosaProfileCompliance.cpp
@@ -329,6 +329,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(SliceShape)
POPULATE_PROFILE_INFO_SKIP(SubShape)
diff --git a/mlir/lib/Dialect/Tosa/Transforms/TosaValidation.cpp b/mlir/lib/Dialect/Tosa/Transforms/TosaValidation.cpp
index 4b9c8b2030a49..99eb83ad6d12c 100644
--- a/mlir/lib/Dialect/Tosa/Transforms/TosaValidation.cpp
+++ b/mlir/lib/Dialect/Tosa/Transforms/TosaValidation.cpp
@@ -708,6 +708,7 @@ LogicalResult TosaValidation::levelCheckRanksAndSizes(Operation *op) {
CHECK_RANKS(ConcatShape);
CHECK_RANKS(DivCeilShape);
CHECK_RANKS(DivFloorShape);
+ CHECK_RANKS(ModShape);
CHECK_RANKS(MulShape);
CHECK_RANKS(SliceShape);
CHECK_RANKS(SubShape);
diff --git a/mlir/test/Dialect/Tosa/level_check.mlir b/mlir/test/Dialect/Tosa/level_check.mlir
index ec540bda5e57d..5e681c1ef75c8 100644
--- a/mlir/test/Dialect/Tosa/level_check.mlir
+++ b/mlir/test/Dialect/Tosa/level_check.mlir
@@ -1719,3 +1719,13 @@ func.func @test_concat_shape_invalid_list_size() {
) -> !tosa.shape<0>
return
}
+
+// -----
+
+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 276eac4d6166d..9da17cf4f9552 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 63379ed8d8a4d..1b2403c60fbbd 100644
--- a/mlir/test/Dialect/Tosa/tosa-validation-version-1p1-valid.mlir
+++ b/mlir/test/Dialect/Tosa/tosa-validation-version-1p1-valid.mlir
@@ -166,3 +166,13 @@ func.func @test_dim(%arg0: tensor<1x2x3x4xi32>) -> !tosa.shape<1> {
%0 = tosa.dim %arg0 {axis = 2 : i32} : (tensor<1x2x3x4xi32>) -> !tosa.shape<1>
return %0 : !tosa.shape<1>
}
+
+// -----
+
+// 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>
+}
>From 8b572de546ae7583448b4a64b1d1ba2649578759 Mon Sep 17 00:00:00 2001
From: Luke Hutton <luke.hutton at arm.com>
Date: Thu, 8 Jan 2026 21:23:26 +0000
Subject: [PATCH 2/2] Add rank mismatch test
Change-Id: I1d8e36adcda719b93739b79e6c2a6bdafb4dbc31
---
mlir/test/Dialect/Tosa/verifier.mlir | 10 ++++++++++
1 file changed, 10 insertions(+)
diff --git a/mlir/test/Dialect/Tosa/verifier.mlir b/mlir/test/Dialect/Tosa/verifier.mlir
index a51ed4f09400f..e444664cf2b93 100644
--- a/mlir/test/Dialect/Tosa/verifier.mlir
+++ b/mlir/test/Dialect/Tosa/verifier.mlir
@@ -1325,3 +1325,13 @@ func.func @test_slice_shape_incorrect_output_size() -> !tosa.shape<4> {
%slice = tosa.slice_shape %shape, %start, %size : (!tosa.shape<6>, tensor<1xi32>, tensor<1xi32>) -> !tosa.shape<4>
return %slice : !tosa.shape<4>
}
+
+// -----
+
+func.func @test_mod_shape_input1_input2_rank_mismatch() -> !tosa.shape<6> {
+ %a = tosa.const_shape {values = dense<[1, 2, 3, 4, 5, 6]> : tensor<6xindex>} : () -> !tosa.shape<6>
+ %b = tosa.const_shape {values = dense<[1, 2, 3, 4, 5]> : tensor<5xindex>} : () -> !tosa.shape<5>
+ // expected-error at +1 {{'tosa.mod_shape' op operands don't have matching ranks}}
+ %c = tosa.mod_shape %a, %b : (!tosa.shape<6>, !tosa.shape<5>) -> !tosa.shape<6>
+ return %c : !tosa.shape<6>
+}
More information about the Mlir-commits
mailing list