[Mlir-commits] [mlir] [mlir][mesh] Add collective communication operations (PR #71960)

[Boian Petkantchin]

================
@@ -0,0 +1,72 @@
+// RUN: mlir-opt --canonicalize %s | FileCheck %s
+
+mesh.cluster @mesh0(rank = 2, dim_sizes = [2, 4])
+
+// CHECK-LABEL: func @all_reduce_mesh_axes
+func.func @all_reduce_mesh_axes(
+    %arg0 : tensor<4xf32>) -> tensor<4xf64> {
+// CHECK: mesh_axes = array<i16: 0, 1>
+  %0 = mesh.all_reduce %arg0 {
+    mesh = @mesh0, mesh_axes = array<i16: 1, 0, 0>, reduction = #mesh.partial<sum>
+    } : tensor<4xf32> -> tensor<4xf64>
+  return %0 : tensor<4xf64>
+}
+
+// CHECK-LABEL: func @all_reduce_empty_mesh_axes_and_default_reduction
----------------
sogartar wrote:

An empty `mesh_axes` should mean each device in its own group. I think it should be valid. It would have to be optimized out. If we have the rewrite for the empty case then all other passes can benefit and not that it should be a no op.

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