[Mlir-commits] [mlir] mlir::mesh::shardingOp adding shard-size control (PR #98145)

[Boian Petkantchin]

================
@@ -151,10 +178,86 @@ func.func @incomplete_sharding(
   %arg0: tensor<8x16xf32>
 // CHECK-SAME: -> tensor<4x16xf32> {
 ) -> tensor<8x16xf32> {
-  %0 = mesh.shard %arg0 to <@mesh_1d, [[0]]> annotate_for_users : tensor<8x16xf32>
+  %s0 = mesh.sharding @mesh_1d, [[0]] : !mesh.sharding
+  %0 = mesh.shard %arg0 to %s0  annotate_for_users : tensor<8x16xf32>
   // CHECK: %[[RES:.*]] = tosa.sigmoid %[[ARG]] : (tensor<4x16xf32>) -> tensor<4x16xf32>
   %1 = tosa.sigmoid %0 : (tensor<8x16xf32>) -> tensor<8x16xf32>
-  %2 = mesh.shard %1 to <@mesh_1d, [[0]]> : tensor<8x16xf32>
+  %s2 = mesh.sharding @mesh_1d, [[0]] : !mesh.sharding
+  %2 = mesh.shard %1 to %s2  : tensor<8x16xf32>
   // CHECK: return %[[RES]] : tensor<4x16xf32>
   return %2 : tensor<8x16xf32>
 }
+
+mesh.mesh @mesh_1d_4(shape = 4)
+// CHECK-LABEL: func @update_halo_static
+func.func @update_halo_static(
+  // CHECK-SAME: %[[IN1:[A-Za-z0-9_]+]]: tensor<11x16xi8>
+  %in1: tensor<32x16xi8>
+  // CHECK-SAME: -> tensor<11x16xi8> {
+) -> tensor<32x16xi8> {
+  // CHECK: %[[RES:.*]] = mesh.update_halo %[[IN1]] on @mesh_1d_4  mesh_axes = [0] halo_sizes = [2, 1] : (tensor<11x16xi8>) -> tensor<11x16xi8>
----------------
sogartar wrote:

I don't understand why this function does need to do anything else then return the same tensor.
The input and output shardings are the same.

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