[Mlir-commits] [mlir] 01d70ed - [mlir][shard] Bounds-check the axis attributes in the collective verifiers (#219691)

[llvmlistbot at llvm.org]

Author: Alessandro Potenza
Date: 2026-09-01T15:01:01+02:00
New Revision: 01d70ed4fe59cabf2390bd5931bb6a2f7a34a161

URL: https://github.com/llvm/llvm-project/commit/01d70ed4fe59cabf2390bd5931bb6a2f7a34a161
DIFF: https://github.com/llvm/llvm-project/commit/01d70ed4fe59cabf2390bd5931bb6a2f7a34a161.diff

LOG: [mlir][shard] Bounds-check the axis attributes in the collective verifiers (#219691)

Fixes #218212.

Added: 
    

Modified: 
    mlir/lib/Dialect/Shard/IR/ShardOps.cpp
    mlir/test/Dialect/Shard/invalid.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Dialect/Shard/IR/ShardOps.cpp b/mlir/lib/Dialect/Shard/IR/ShardOps.cpp
index ff790a0bf961d..a4206d6717687 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(
@@ -1127,9 +1138,16 @@ 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())
+           << tensorAxisAttrName << " " << tensorAxis
+           << " is out of bounds [0, " << operandRank << ").";
+  }
   for (int64_t axis = 0; axis < operandType.getRank(); ++axis) {
     if (axis != tensorAxis) {
       if (failed(verifyDimensionCompatibility(
@@ -1241,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,
@@ -1422,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,
@@ -1452,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 c92932a725d1c..b2f19e60e7fdf 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 {{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>
+  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 {{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>
+  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 {{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>
+}
+
+// -----
+
+shard.grid @grid0(shape = 2x2x4)
+
+func.func @reduce_scatter_invalid_scatter_dim(
+    %arg0 : tensor<3x4xf32>) -> tensor<3x4xf64> {
+  // 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>
+  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>
+}