checking for non-decreasing dims sizes in sharding

Author: fschlimb

mlir/lib/Dialect/Mesh/IR/MeshOps.cpp

@@ -198,11 +198,12 @@ static void shardShape(const InShape &inShape, const MeshShape &meshShape,
             llvm::adl_begin(outShape));
 
   if (!shardedDimsOffsets.empty()) {
+    auto isDynShape = ShapedType::isDynamicShape(meshShape);
     uint64_t pos = 1;
     for (auto [tensorAxis, innerSplitAxes] : llvm::enumerate(splitAxes)) {
       if (!innerSplitAxes.empty()) {
         auto sz = shardedDimsOffsets[pos];
-        bool same = !ShapedType::isDynamicShape(meshShape);
+        bool same = !isDynShape;
         if (same) {
           // Find sharded dims in shardedDimsOffsets with same static size on
           // all devices. Use kDynamic for dimensions with dynamic or
@@ -218,7 +219,7 @@ static void shardShape(const InShape &inShape, const MeshShape &meshShape,
               break;
             }
           }
-          pos += numShards;
+          pos += numShards + 1;
         }
         outShape[tensorAxis] = same ? sz : ShapedType::kDynamic;
       }
@@ -544,6 +545,34 @@ LogicalResult ShardingOp::verifySymbolUses(SymbolTableCollection &symbolTable) {
     return emitError() << "sharded dims offsets are not allowed for "
                           "devices meshes with dynamic shape.";
   }
+
+  auto shardedDimsOffsets = getStaticShardedDimsOffsets();
+  if (!shardedDimsOffsets.empty()) {
+    auto meshShape = mesh.value().getShape();
+    assert(!ShapedType::isDynamicShape(meshShape));
+    uint64_t pos = 0;
+    for (auto [tensorAxis, innerSplitAxes] : llvm::enumerate(getSplitAxes())) {
+      if (!innerSplitAxes.empty()) {
+        int64_t numShards = 0, off = 0;
+        for (auto i : innerSplitAxes.asArrayRef()) {
+          numShards += meshShape[i];
+        }
+        for (int64_t i = 0; i <= numShards; ++i) {
+          if (shardedDimsOffsets.size() <= pos + i) {
+            return emitError() << "sharded dims offsets has wrong size.";
+          }
+          if (!ShapedType::isDynamic(shardedDimsOffsets[pos + i])) {
+            if (shardedDimsOffsets[pos + i] < off) {
+              return emitError()
+                     << "sharded dims offsets must be non-decreasing.";
+            }
+            off = shardedDimsOffsets[pos + i];
+          }
+        }
+        pos += numShards + 1;
+      }
+    }
+  }
   return success();
 }
 

mlir/test/Dialect/Mesh/canonicalization.mlir

@@ -204,7 +204,7 @@ func.func @test_halo_sizes() -> !mesh.sharding {
 // CHECK-LABEL: func @test_shard_offs
 func.func @test_shard_offs() -> !mesh.sharding {
   %c2_i64 = arith.constant 2 : i64
-  // CHECK mesh.sharding @mesh4x4 split_axes = [[0], [1]] sharded_dims_offsets = [0, 1, 2, 0, 2, 22] : !mesh.sharding
-  %sharding = mesh.sharding @mesh4x4 split_axes = [[0], [1]] sharded_dims_offsets = [0, 1, %c2_i64, 0, %c2_i64, 22] : !mesh.sharding
+  // CHECK mesh.sharding @mesh4x4 split_axes = [[0], [1]] sharded_dims_offsets = [0, 1, 2, 3, 4, 0, 2, 3, 4, 22] : !mesh.sharding
+  %sharding = mesh.sharding @mesh4x4 split_axes = [[0], [1]] sharded_dims_offsets = [0, 1, %c2_i64, 3, 4, 0, %c2_i64, 3, 4, 22] : !mesh.sharding
   return %sharding : !mesh.sharding
 }

mlir/test/Dialect/Mesh/invalid.mlir

@@ -107,6 +107,26 @@ func.func @sharding_dyn_mesh_and_sizes(%arg0 : tensor<4x8xf32>) {
 
 // -----
 
+mesh.mesh @mesh0(shape = 2x4)
+func.func @sharding_sizes_count(%arg0 : tensor<4x8xf32>) {
+  // expected-error@+1 {{sharded dims offsets has wrong size}}
+  %s = mesh.sharding @mesh0 split_axes = [[0], [1]] sharded_dims_offsets = [0, 2, 4, 0, 2, 4, 6] : !mesh.sharding
+  %0 = mesh.shard %arg0 to %s : tensor<4x8xf32>
+  return
+}
+
+// -----
+
+mesh.mesh @mesh0(shape = 4)
+func.func @sharding_sizes_decreasing(%arg0 : tensor<4x8xf32>) {
+  // expected-error@+1 {{sharded dims offsets must be non-decreasing}}
+  %s = mesh.sharding @mesh0 split_axes = [[0]] sharded_dims_offsets = [0, 2, 3, 2] : !mesh.sharding
+  %0 = mesh.shard %arg0 to %s : tensor<4x8xf32>
+  return
+}
+
+// -----
+
 mesh.mesh @mesh0(shape = 2x4)
 
 func.func @mesh_shape_mesh_axis_out_of_bounds() -> (index, index) {