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

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir

Author: Alessandro Potenza (alepot55)

<details>
<summary>Changes</summary>

`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.

---
Full diff: https://github.com/llvm/llvm-project/pull/219691.diff


2 Files Affected:

- (modified) mlir/lib/Dialect/Shard/IR/ShardOps.cpp (+17) 
- (modified) mlir/test/Dialect/Shard/invalid.mlir (+77) 


``````````diff
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>
+}

``````````

</details>


https://github.com/llvm/llvm-project/pull/219691