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

Alessandro Potenza llvmlistbot at llvm.org
Mon Aug 31 01:49:24 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 7caa947dc9bcd4538a5ed1015380707cb067f814 Mon Sep 17 00:00:00 2001
From: Alessandro Potenza <ap.alessandro.potenza at gmail.com>
Date: Mon, 31 Aug 2026 10:49:09 +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,
because the attribute is `slice_axis` for all_slice and `scatter_dim` for
scatter and reduce_scatter. Pass the name in and print it verbatim, and do the
same for all_to_all's `split_axis` and `concat_axis`.
---
 .../AMDGPU/MCTargetDesc/AMDGPUInstPrinter.h   |  1 +
 llvm/unittests/IR/VPIntrinsicTest.cpp         | 82 +++++++++++--------
 mlir/lib/Dialect/Shard/IR/ShardOps.cpp        | 25 +++---
 mlir/test/Dialect/Shard/invalid.mlir          | 12 +--
 4 files changed, 70 insertions(+), 50 deletions(-)

diff --git a/llvm/lib/Target/AMDGPU/MCTargetDesc/AMDGPUInstPrinter.h b/llvm/lib/Target/AMDGPU/MCTargetDesc/AMDGPUInstPrinter.h
index 5e9ebc6716c7f..2ee5b42038287 100644
--- a/llvm/lib/Target/AMDGPU/MCTargetDesc/AMDGPUInstPrinter.h
+++ b/llvm/lib/Target/AMDGPU/MCTargetDesc/AMDGPUInstPrinter.h
@@ -202,6 +202,7 @@ class AMDGPUInstPrinter : public MCInstPrinter {
                          StringRef Asm, StringRef Default = "");
   static void printIfSet(const MCInst *MI, unsigned OpNo, raw_ostream &O,
                          char Asm);
+
 protected:
   void printAbs(const MCInst *MI, unsigned OpNo, const MCSubtargetInfo &STI,
                 raw_ostream &O);
diff --git a/llvm/unittests/IR/VPIntrinsicTest.cpp b/llvm/unittests/IR/VPIntrinsicTest.cpp
index 66e12a5fbac29..8fd00aea0ab98 100644
--- a/llvm/unittests/IR/VPIntrinsicTest.cpp
+++ b/llvm/unittests/IR/VPIntrinsicTest.cpp
@@ -40,10 +40,9 @@ class VPIntrinsicTest : public testing::Test {
   SMDiagnostic Err;
 
   std::unique_ptr<Module> createVPDeclarationModule() {
-    const char *BinaryIntOpcodes[] = {"add",  "sub",  "mul", "sdiv", "srem",
-                                      "udiv", "urem", "and", "xor",  "or",
-                                      "ashr", "lshr", "shl", "smin", "smax",
-                                      "umin", "umax"};
+    const char *BinaryIntOpcodes[] = {
+        "add", "sub",  "mul",  "sdiv", "srem", "udiv", "urem", "and", "xor",
+        "or",  "ashr", "lshr", "shl",  "smin", "smax", "umin", "umax"};
     std::stringstream Str;
     for (const char *BinaryIntOpcode : BinaryIntOpcodes)
       Str << " declare <8 x i32> @llvm.vp." << BinaryIntOpcode
@@ -236,34 +235,53 @@ TEST_F(VPIntrinsicTest, CanIgnoreVectorLength) {
   LLVMContext C;
   SMDiagnostic Err;
 
-  std::unique_ptr<Module> M =
-      parseAssemblyString(
-"declare <256 x i64> @llvm.vp.mul.v256i64(<256 x i64>, <256 x i64>, <256 x i1>, i32)"
-"declare <vscale x 2 x i64> @llvm.vp.mul.nxv2i64(<vscale x 2 x i64>, <vscale x 2 x i64>, <vscale x 2 x i1>, i32)"
-"declare <vscale x 1 x i64> @llvm.vp.mul.nxv1i64(<vscale x 1 x i64>, <vscale x 1 x i64>, <vscale x 1 x i1>, i32)"
-"declare i32 @llvm.vscale.i32()"
-"define void @test_static_vlen( "
-"      <256 x i64> %i0, <vscale x 2 x i64> %si0x2, <vscale x 1 x i64> %si0x1,"
-"      <256 x i64> %i1, <vscale x 2 x i64> %si1x2, <vscale x 1 x i64> %si1x1,"
-"      <256 x i1> %m, <vscale x 2 x i1> %smx2, <vscale x 1 x i1> %smx1, i32 %vl) { "
-"  %r0 = call <256 x i64> @llvm.vp.mul.v256i64(<256 x i64> %i0, <256 x i64> %i1, <256 x i1> %m, i32 %vl)"
-"  %r1 = call <256 x i64> @llvm.vp.mul.v256i64(<256 x i64> %i0, <256 x i64> %i1, <256 x i1> %m, i32 256)"
-"  %r2 = call <256 x i64> @llvm.vp.mul.v256i64(<256 x i64> %i0, <256 x i64> %i1, <256 x i1> %m, i32 0)"
-"  %r3 = call <256 x i64> @llvm.vp.mul.v256i64(<256 x i64> %i0, <256 x i64> %i1, <256 x i1> %m, i32 7)"
-"  %r4 = call <256 x i64> @llvm.vp.mul.v256i64(<256 x i64> %i0, <256 x i64> %i1, <256 x i1> %m, i32 123)"
-"  %vs = call i32 @llvm.vscale.i32()"
-"  %vs.x2 = mul i32 %vs, 2"
-"  %r5 = call <vscale x 2 x i64> @llvm.vp.mul.nxv2i64(<vscale x 2 x i64> %si0x2, <vscale x 2 x i64> %si1x2, <vscale x 2 x i1> %smx2, i32 %vs.x2)"
-"  %r6 = call <vscale x 2 x i64> @llvm.vp.mul.nxv2i64(<vscale x 2 x i64> %si0x2, <vscale x 2 x i64> %si1x2, <vscale x 2 x i1> %smx2, i32 %vs)"
-"  %r7 = call <vscale x 2 x i64> @llvm.vp.mul.nxv2i64(<vscale x 2 x i64> %si0x2, <vscale x 2 x i64> %si1x2, <vscale x 2 x i1> %smx2, i32 99999)"
-"  %r8 = call <vscale x 1 x i64> @llvm.vp.mul.nxv1i64(<vscale x 1 x i64> %si0x1, <vscale x 1 x i64> %si1x1, <vscale x 1 x i1> %smx1, i32 %vs)"
-"  %r9 = call <vscale x 1 x i64> @llvm.vp.mul.nxv1i64(<vscale x 1 x i64> %si0x1, <vscale x 1 x i64> %si1x1, <vscale x 1 x i1> %smx1, i32 1)"
-"  %r10 = call <vscale x 1 x i64> @llvm.vp.mul.nxv1i64(<vscale x 1 x i64> %si0x1, <vscale x 1 x i64> %si1x1, <vscale x 1 x i1> %smx1, i32 %vs.x2)"
-"  %vs.wat = add i32 %vs, 2"
-"  %r11 = call <vscale x 2 x i64> @llvm.vp.mul.nxv2i64(<vscale x 2 x i64> %si0x2, <vscale x 2 x i64> %si1x2, <vscale x 2 x i1> %smx2, i32 %vs.wat)"
-"  ret void "
-"}",
-          Err, C);
+  std::unique_ptr<Module> M = parseAssemblyString(
+      "declare <256 x i64> @llvm.vp.mul.v256i64(<256 x i64>, <256 x i64>, <256 "
+      "x i1>, i32)"
+      "declare <vscale x 2 x i64> @llvm.vp.mul.nxv2i64(<vscale x 2 x i64>, "
+      "<vscale x 2 x i64>, <vscale x 2 x i1>, i32)"
+      "declare <vscale x 1 x i64> @llvm.vp.mul.nxv1i64(<vscale x 1 x i64>, "
+      "<vscale x 1 x i64>, <vscale x 1 x i1>, i32)"
+      "declare i32 @llvm.vscale.i32()"
+      "define void @test_static_vlen( "
+      "      <256 x i64> %i0, <vscale x 2 x i64> %si0x2, <vscale x 1 x i64> "
+      "%si0x1,"
+      "      <256 x i64> %i1, <vscale x 2 x i64> %si1x2, <vscale x 1 x i64> "
+      "%si1x1,"
+      "      <256 x i1> %m, <vscale x 2 x i1> %smx2, <vscale x 1 x i1> %smx1, "
+      "i32 %vl) { "
+      "  %r0 = call <256 x i64> @llvm.vp.mul.v256i64(<256 x i64> %i0, <256 x "
+      "i64> %i1, <256 x i1> %m, i32 %vl)"
+      "  %r1 = call <256 x i64> @llvm.vp.mul.v256i64(<256 x i64> %i0, <256 x "
+      "i64> %i1, <256 x i1> %m, i32 256)"
+      "  %r2 = call <256 x i64> @llvm.vp.mul.v256i64(<256 x i64> %i0, <256 x "
+      "i64> %i1, <256 x i1> %m, i32 0)"
+      "  %r3 = call <256 x i64> @llvm.vp.mul.v256i64(<256 x i64> %i0, <256 x "
+      "i64> %i1, <256 x i1> %m, i32 7)"
+      "  %r4 = call <256 x i64> @llvm.vp.mul.v256i64(<256 x i64> %i0, <256 x "
+      "i64> %i1, <256 x i1> %m, i32 123)"
+      "  %vs = call i32 @llvm.vscale.i32()"
+      "  %vs.x2 = mul i32 %vs, 2"
+      "  %r5 = call <vscale x 2 x i64> @llvm.vp.mul.nxv2i64(<vscale x 2 x i64> "
+      "%si0x2, <vscale x 2 x i64> %si1x2, <vscale x 2 x i1> %smx2, i32 %vs.x2)"
+      "  %r6 = call <vscale x 2 x i64> @llvm.vp.mul.nxv2i64(<vscale x 2 x i64> "
+      "%si0x2, <vscale x 2 x i64> %si1x2, <vscale x 2 x i1> %smx2, i32 %vs)"
+      "  %r7 = call <vscale x 2 x i64> @llvm.vp.mul.nxv2i64(<vscale x 2 x i64> "
+      "%si0x2, <vscale x 2 x i64> %si1x2, <vscale x 2 x i1> %smx2, i32 99999)"
+      "  %r8 = call <vscale x 1 x i64> @llvm.vp.mul.nxv1i64(<vscale x 1 x i64> "
+      "%si0x1, <vscale x 1 x i64> %si1x1, <vscale x 1 x i1> %smx1, i32 %vs)"
+      "  %r9 = call <vscale x 1 x i64> @llvm.vp.mul.nxv1i64(<vscale x 1 x i64> "
+      "%si0x1, <vscale x 1 x i64> %si1x1, <vscale x 1 x i1> %smx1, i32 1)"
+      "  %r10 = call <vscale x 1 x i64> @llvm.vp.mul.nxv1i64(<vscale x 1 x "
+      "i64> %si0x1, <vscale x 1 x i64> %si1x1, <vscale x 1 x i1> %smx1, i32 "
+      "%vs.x2)"
+      "  %vs.wat = add i32 %vs, 2"
+      "  %r11 = call <vscale x 2 x i64> @llvm.vp.mul.nxv2i64(<vscale x 2 x "
+      "i64> %si0x2, <vscale x 2 x i64> %si1x2, <vscale x 2 x i1> %smx2, i32 "
+      "%vs.wat)"
+      "  ret void "
+      "}",
+      Err, C);
 
   auto *F = M->getFunction("test_static_vlen");
   assert(F);
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