[Mlir-commits] [mlir] [mlir][shard] Bounds-check the axis attributes in the collective verifiers (PR #219691)
Alessandro Potenza
llvmlistbot at llvm.org
Mon Aug 31 01:52:47 PDT 2026
https://github.com/alepot55 updated https://github.com/llvm/llvm-project/pull/219691
>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 1/2] [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>
+}
>From 224c6722c4fb60170a209319816759d26ce30ef3 Mon Sep 17 00:00:00 2001
From: Alessandro Potenza <ap.alessandro.potenza at gmail.com>
Date: Mon, 31 Aug 2026 10:52:34 +0200
Subject: [PATCH 2/2] Name the attribute in the out-of-bounds messages
Review feedback: "Tensor axis" did not correspond to anything the user wrote.
The attribute is `slice_axis` for all_slice and `scatter_dim` for scatter and
reduce_scatter, so one generic word stood for two different names. Pass the
name in and print it verbatim, and do the same for all_to_all's `split_axis`
and `concat_axis`.
---
mlir/lib/Dialect/Shard/IR/ShardOps.cpp | 25 +++++++++++++------------
mlir/test/Dialect/Shard/invalid.mlir | 12 ++++++------
2 files changed, 19 insertions(+), 18 deletions(-)
diff --git a/mlir/lib/Dialect/Shard/IR/ShardOps.cpp b/mlir/lib/Dialect/Shard/IR/ShardOps.cpp
index bce58f84205b2..a4206d6717687 100644
--- a/mlir/lib/Dialect/Shard/IR/ShardOps.cpp
+++ b/mlir/lib/Dialect/Shard/IR/ShardOps.cpp
@@ -1088,12 +1088,12 @@ static LogicalResult verifyAllToAllOperandAndResultShape(
auto operandRank = operandType.getRank();
if (splitAxis < 0 || splitAxis >= operandRank) {
return emitError(result.getLoc())
- << "Split axis " << splitAxis << " is out of bounds [0, "
+ << "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, "
+ << "concat_axis " << concatAxis << " is out of bounds [0, "
<< operandRank << ").";
}
for (int64_t axis = 0; axis < operandType.getRank(); ++axis) {
@@ -1138,14 +1138,15 @@ static LogicalResult verifyAllToAllOperandAndResultShape(
static LogicalResult verifyScatterOrSliceOperandAndResultShape(
Value operand, Value result, int64_t tensorAxis,
- ArrayRef<GridAxis> gridAxes, ArrayRef<int64_t> gridShape) {
+ StringRef tensorAxisAttrName, 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 << ").";
+ << tensorAxisAttrName << " " << tensorAxis
+ << " is out of bounds [0, " << operandRank << ").";
}
for (int64_t axis = 0; axis < operandType.getRank(); ++axis) {
if (axis != tensorAxis) {
@@ -1258,8 +1259,8 @@ LogicalResult AllSliceOp::verifySymbolUses(SymbolTableCollection &symbolTable) {
return failure();
}
return verifyScatterOrSliceOperandAndResultShape(
- getOperand(), getResult(), getSliceAxis().getSExtValue(), getGridAxes(),
- grid.value().getShape());
+ getOperand(), getResult(), getSliceAxis().getSExtValue(), "slice_axis",
+ getGridAxes(), grid.value().getShape());
}
void AllSliceOp::getCanonicalizationPatterns(RewritePatternSet &patterns,
@@ -1439,8 +1440,8 @@ ReduceScatterOp::verifySymbolUses(SymbolTableCollection &symbolTable) {
}
return verifyScatterOrSliceOperandAndResultShape(
- getOperand(), getResult(), getScatterDim().getSExtValue(), getGridAxes(),
- grid.value().getShape());
+ getOperand(), getResult(), getScatterDim().getSExtValue(), "scatter_dim",
+ getGridAxes(), grid.value().getShape());
}
void ReduceScatterOp::getCanonicalizationPatterns(RewritePatternSet &patterns,
@@ -1469,9 +1470,9 @@ LogicalResult ScatterOp::verifySymbolUses(SymbolTableCollection &symbolTable) {
}
auto scatterDim = getScatterDim().getSExtValue();
- return verifyScatterOrSliceOperandAndResultShape(getInput(), getResult(),
- scatterDim, getGridAxes(),
- grid.value().getShape());
+ return verifyScatterOrSliceOperandAndResultShape(
+ getInput(), getResult(), scatterDim, "scatter_dim", getGridAxes(),
+ grid.value().getShape());
}
void ScatterOp::getCanonicalizationPatterns(RewritePatternSet &patterns,
diff --git a/mlir/test/Dialect/Shard/invalid.mlir b/mlir/test/Dialect/Shard/invalid.mlir
index ddb889d2352a3..b2f19e60e7fdf 100644
--- a/mlir/test/Dialect/Shard/invalid.mlir
+++ b/mlir/test/Dialect/Shard/invalid.mlir
@@ -934,7 +934,7 @@ 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).}}
+ // expected-error at +1 {{scatter_dim 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>
@@ -947,7 +947,7 @@ 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).}}
+ // expected-error at +1 {{scatter_dim -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>
@@ -960,7 +960,7 @@ 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).}}
+ // expected-error at +1 {{slice_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>
@@ -972,7 +972,7 @@ 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).}}
+ // expected-error at +1 {{scatter_dim 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>
@@ -985,7 +985,7 @@ 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).}}
+ // 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>
@@ -998,7 +998,7 @@ 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).}}
+ // 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>
More information about the Mlir-commits
mailing list