[Mlir-commits] [mlir] [mlir][shard] Bounds-check the axis attributes in the collective verifiers (PR #219691)
Alessandro Potenza
llvmlistbot at llvm.org
Sat Aug 29 08:34:40 PDT 2026
https://github.com/alepot55 created https://github.com/llvm/llvm-project/pull/219691
`verifyGatherOperandAndResultShape` rejects an out-of-range axis before using it. Its two siblings do not: they pass the attribute straight to `ShapedType::getDimSize`, which asserts in an assertions build and reads past the end of the shape otherwise. The verifier is the thing that should be catching this.
Four ops reach the unchecked path: `shard.scatter`, `shard.all_slice` and `shard.reduce_scatter` through `verifyScatterOrSliceOperandAndResultShape`, and `shard.all_to_all` through `verifyAllToAllOperandAndResultShape`.
Built `mlir-opt` with `-DLLVM_ENABLE_ASSERTIONS=ON` and ran one reduced case per op, before and after:
| case | before | after |
|---|---|---|
| `shard.scatter`, `scatter_dim = 5` | `rc=134`, `Assertion 'idx < getRank()' failed` | `Tensor axis 5 is out of bounds [0, 2).` |
| `shard.all_slice`, `slice_axis = 5` | `rc=134`, same assertion | `Tensor axis 5 is out of bounds [0, 2).` |
| `shard.reduce_scatter`, `scatter_dim = 5` | `rc=134`, same assertion | `Tensor axis 5 is out of bounds [0, 2).` |
| `shard.all_to_all`, `split_axis = 5` | `rc=134`, same assertion | `Split axis 5 is out of bounds [0, 2).` |
| `shard.gather`, `gather_axis = 5` (control, untouched) | `Gather axis 5 is out of bounds [0, 2).` | unchanged |
Logs: [before](https://github.com/alepot55/llvm-project/actions/runs/33248380586), [after](https://github.com/alepot55/llvm-project/actions/runs/33256402112).
The bound is the operand rank, which is the first out-of-range read and the bound the existing per-axis loop already uses.
One note on the reduced cases, since it cost me a round: with an out-of-range axis that per-axis loop compares *every* axis, so operand and result must have the same shape. Give the result the shape a valid collective would produce and an earlier dimension-mismatch error fires first and the axis is never reached.
Tests: six cases in `invalid.mlir`, one per path plus a negative axis. `ninja check-mlir-dialect-shard` 16/16, `ninja check-mlir` 3928 passed, 615 unsupported, 1 expectedly failed, 0 failed. `git clang-format` against the merge base is clean.
Disclosure per the [AI tool policy](https://llvm.org/docs/AIToolPolicy.html): AI assistance was used to find this and to prepare the patch. I understand the change and can defend it in review.
Fixes #218212.
>From a959bc171fea7777d8e54b1ad02564317ecc2681 Mon Sep 17 00:00:00 2001
From: Alessandro Potenza <ap.alessandro.potenza at gmail.com>
Date: Sat, 29 Aug 2026 15:59:16 +0200
Subject: [PATCH] [mlir][shard] Bounds-check the axis attributes in the
collective verifiers
verifyGatherOperandAndResultShape rejects an out-of-range axis before using
it. verifyAllToAllOperandAndResultShape and
verifyScatterOrSliceOperandAndResultShape do not, and pass the attribute
straight to ShapedType::getDimSize, which asserts in an assertions build and
reads past the end of the shape otherwise. The verifier is what should be
catching this, so it is a poor place to go out of bounds.
Reached from shard.scatter, shard.all_slice, shard.reduce_scatter and
shard.all_to_all. All four abort on an out-of-range axis today, while
shard.gather with the same axis reports it cleanly.
Add the same check to both, wording it like gather's, and cover each path in
invalid.mlir.
Assisted-by: Claude (Anthropic)
---
mlir/lib/Dialect/Shard/IR/ShardOps.cpp | 17 ++++++
mlir/test/Dialect/Shard/invalid.mlir | 77 ++++++++++++++++++++++++++
2 files changed, 94 insertions(+)
diff --git a/mlir/lib/Dialect/Shard/IR/ShardOps.cpp b/mlir/lib/Dialect/Shard/IR/ShardOps.cpp
index ff790a0bf961d..bce58f84205b2 100644
--- a/mlir/lib/Dialect/Shard/IR/ShardOps.cpp
+++ b/mlir/lib/Dialect/Shard/IR/ShardOps.cpp
@@ -1085,6 +1085,17 @@ static LogicalResult verifyAllToAllOperandAndResultShape(
ArrayRef<GridAxis> gridAxes, ArrayRef<int64_t> gridShape) {
ShapedType operandType = cast<ShapedType>(operand.getType());
ShapedType resultType = cast<ShapedType>(result.getType());
+ auto operandRank = operandType.getRank();
+ if (splitAxis < 0 || splitAxis >= operandRank) {
+ return emitError(result.getLoc())
+ << "Split axis " << splitAxis << " is out of bounds [0, "
+ << operandRank << ").";
+ }
+ if (concatAxis < 0 || concatAxis >= operandRank) {
+ return emitError(result.getLoc())
+ << "Concat axis " << concatAxis << " is out of bounds [0, "
+ << operandRank << ").";
+ }
for (int64_t axis = 0; axis < operandType.getRank(); ++axis) {
if ((axis != splitAxis && axis != concatAxis) || splitAxis == concatAxis) {
if (failed(verifyDimensionCompatibility(
@@ -1130,6 +1141,12 @@ static LogicalResult verifyScatterOrSliceOperandAndResultShape(
ArrayRef<GridAxis> gridAxes, ArrayRef<int64_t> gridShape) {
ShapedType operandType = cast<ShapedType>(operand.getType());
ShapedType resultType = cast<ShapedType>(result.getType());
+ auto operandRank = operandType.getRank();
+ if (tensorAxis < 0 || tensorAxis >= operandRank) {
+ return emitError(result.getLoc())
+ << "Tensor axis " << tensorAxis << " is out of bounds [0, "
+ << operandRank << ").";
+ }
for (int64_t axis = 0; axis < operandType.getRank(); ++axis) {
if (axis != tensorAxis) {
if (failed(verifyDimensionCompatibility(
diff --git a/mlir/test/Dialect/Shard/invalid.mlir b/mlir/test/Dialect/Shard/invalid.mlir
index c92932a725d1c..ddb889d2352a3 100644
--- a/mlir/test/Dialect/Shard/invalid.mlir
+++ b/mlir/test/Dialect/Shard/invalid.mlir
@@ -927,3 +927,80 @@ func.func @shift_invalid_shift_axis(
: tensor<4xi8> -> tensor<4xi8>
return %0 : tensor<4xi8>
}
+
+// -----
+
+shard.grid @grid0(shape = 2x2x4)
+
+func.func @scatter_invalid_scatter_dim(
+ %arg0 : tensor<3x4xf32>) -> tensor<3x4xf32> {
+ // expected-error at +1 {{Tensor axis 5 is out of bounds [0, 2).}}
+ %0 = shard.scatter %arg0 on @grid0 grid_axes = [2]
+ scatter_dim = 5 root = [1]
+ : (tensor<3x4xf32>) -> tensor<3x4xf32>
+ return %0 : tensor<3x4xf32>
+}
+
+// -----
+
+shard.grid @grid0(shape = 2x2x4)
+
+func.func @scatter_invalid_negative_scatter_dim(
+ %arg0 : tensor<3x4xf32>) -> tensor<3x4xf32> {
+ // expected-error at +1 {{Tensor axis -1 is out of bounds [0, 2).}}
+ %0 = shard.scatter %arg0 on @grid0 grid_axes = [2]
+ scatter_dim = -1 root = [1]
+ : (tensor<3x4xf32>) -> tensor<3x4xf32>
+ return %0 : tensor<3x4xf32>
+}
+
+// -----
+
+shard.grid @grid0(shape = 2x2x4)
+
+func.func @all_slice_invalid_slice_axis(
+ %arg0 : tensor<3x4xf32>) -> tensor<3x4xf32> {
+ // expected-error at +1 {{Tensor axis 5 is out of bounds [0, 2).}}
+ %0 = shard.all_slice %arg0 on @grid0 grid_axes = [2] slice_axis = 5
+ : tensor<3x4xf32> -> tensor<3x4xf32>
+ return %0 : tensor<3x4xf32>
+}
+
+// -----
+
+shard.grid @grid0(shape = 2x2x4)
+
+func.func @reduce_scatter_invalid_scatter_dim(
+ %arg0 : tensor<3x4xf32>) -> tensor<3x4xf64> {
+ // expected-error at +1 {{Tensor axis 5 is out of bounds [0, 2).}}
+ %0 = shard.reduce_scatter %arg0 on @grid0 grid_axes = [2]
+ reduction = max scatter_dim = 5
+ : tensor<3x4xf32> -> tensor<3x4xf64>
+ return %0 : tensor<3x4xf64>
+}
+
+// -----
+
+shard.grid @grid4(shape = 3)
+
+func.func @all_to_all_invalid_split_axis(
+ %arg0 : tensor<3x6xi8>) -> tensor<3x6xi8> {
+ // expected-error at +1 {{Split axis 5 is out of bounds [0, 2).}}
+ %0 = shard.all_to_all %arg0 on @grid4
+ split_axis = 5 concat_axis = 0
+ : tensor<3x6xi8> -> tensor<3x6xi8>
+ return %0 : tensor<3x6xi8>
+}
+
+// -----
+
+shard.grid @grid4(shape = 3)
+
+func.func @all_to_all_invalid_concat_axis(
+ %arg0 : tensor<3x6xi8>) -> tensor<3x6xi8> {
+ // expected-error at +1 {{Concat axis 5 is out of bounds [0, 2).}}
+ %0 = shard.all_to_all %arg0 on @grid4
+ split_axis = 0 concat_axis = 5
+ : tensor<3x6xi8> -> tensor<3x6xi8>
+ return %0 : tensor<3x6xi8>
+}
More information about the Mlir-commits
mailing list