Automerge: [NFC][mlir][mesh,shard] Fixing misnomers in mesh dialect, renaming 'mesh' dialect to 'shard' (#150177)

Dialect to 'shard' (discourse 87053)
  - dialect name mesh -> shard
  - (device) mesh -> (device) grid
  - spmdize -> partition

A lot of diffs, but simple renames only.

@tkarna @yaochengji

Author: fschlimb

mlir/docs/Dialects/Mesh.md

@@ -0,0 +1,92 @@
+# 'shard' Dialect
+
+The 'shard' dialect defines a set of attributes, operations, and interfaces for
+working with tensor sharding and device communication.
+
+It’s inspired by [GSPMD](*General and Scalable Parallelization for ML Computation Graphs*).
+
+Originally, the dialect was called `mesh`, but it was renamed to better reflect
+what it actually does.
+
+[TOC]
+
+## Collective Communication Operations
+
+The 'shard' dialect includes several collective operations that help coordinate
+communication between devices arranged in a grid.
+
+If you’re not already familiar with collective operations, [this Wikipedia
+article](https://en.wikipedia.org/wiki/Collective_operation) is a good starting
+point.
+
+Unlike traditional collectives that are defined in terms of message-passing
+between explicit buffers on each process, the collectives in this dialect work
+at a higher level. They’re defined in terms of how data moves across the
+dimensions of a tensor, and the participating processes are inferred from how
+the tensor is sharded - not specified manually.
+
+### Device Groups
+
+Each collective operation runs within a group of devices. You define groups
+using the `grid` and `grid_axes` attributes, which describe how to slice the
+full device grid into smaller groups.
+
+Devices that have the same coordinates *outside* the listed `grid_axes` belong
+to the same group.
+
+Example: Say your device grid is shaped `2×3×4×5`, and you set
+`grid_axes = [0, 1]`. This splits the grid into groups by fixing axes 2 and 3. You’d get groups like:
+
+```
+{ { (i, j, k, m) | 0 ≤ i < 2, 0 ≤ j < 3 } | 0 ≤ k < 4, 0 ≤ m < 5 }
+```
+
+So the groups are identified by the coordinates `(k, m)`, and devices like
+`(1, 0, 2, 3)` and `(1, 1, 2, 3)` are in the same group. But `(1, 0, 2, 4)`
+is in a different group.
+
+For some collectives (like `all-to-all`), the order of devices in the group
+matters. The device order is based on the order of axes in `grid_axes`, from
+outermost to innermost.
+
+Example: If `grid_axes = [3, 1]`, then device `(i, 1, k, 0)` comes before
+`(i, 0, k, 1)` and `(i, 2, k, 0)`.
+
+### In-group Devices
+
+Some operations (like `broadcast`, `scatter`, and `send`) refer to a specific
+device within each group. These in-group devices are identified using their
+coordinates over the axes listed in `grid_axes`.
+
+Example: In a 3D grid with `grid_axes = [0, 2]`, an in-group device is specified
+as `(i, j)`. If a group is fixed at coordinate `g` on axis 1, then the full
+device index would be `(i, g, j)`.
+
+### Purity and Execution Model
+
+Collective operations involve all devices in a group (e.g. `all-gather`,
+`all-to-all`) and are considered pure. Operations like `send` and `recv` are not
+collective and are not pure.
+
+The execution model assumes SPMD (Single Program, Multiple Data):
+
+* Every process runs the same program.
+* At any collective operation, all processes are in sync.
+
+This means compiler optimizations must treat collective ops carefully. For
+example, if a collective is removed during optimization, it must be removed from
+*every* path and *every* process that would have participated - otherwise, you’ll
+get undefined behavior at runtime.
+
+Marking these ops as pure also helps with standard compiler passes like dead
+code elimination and common subexpression elimination. It ensures that when the
+program is executed, all devices hit the same line of code at the same time
+during collectives and so avoid dead-locks.
+
+## Operations
+
+[include "Dialects/ShardOps.md"]
+
+## Attributes
+
+[include "Dialects/ShardAttrs.md"]

mlir/docs/Dialects/Shard.md

@@ -72,9 +72,9 @@ This document describes the available MLIR passes and their contracts.
 
 [include "MemRefPasses.md"]
 
-## 'mesh' Dialect Passes
+## 'shard' Dialect Passes
 
-[include "MeshPasses.md"]
+[include "ShardPasses.md"]
 
 ## 'ml\_program' Dialect Passes
 

mlir/docs/Passes.md

@@ -52,7 +52,6 @@
 #include "mlir/Conversion/MemRefToEmitC/MemRefToEmitCPass.h"
 #include "mlir/Conversion/MemRefToLLVM/MemRefToLLVM.h"
 #include "mlir/Conversion/MemRefToSPIRV/MemRefToSPIRVPass.h"
-#include "mlir/Conversion/MeshToMPI/MeshToMPI.h"
 #include "mlir/Conversion/NVGPUToNVVM/NVGPUToNVVM.h"
 #include "mlir/Conversion/NVVMToLLVM/NVVMToLLVM.h"
 #include "mlir/Conversion/OpenACCToSCF/ConvertOpenACCToSCF.h"
@@ -66,6 +65,7 @@
 #include "mlir/Conversion/SCFToSPIRV/SCFToSPIRVPass.h"
 #include "mlir/Conversion/SPIRVToLLVM/SPIRVToLLVMPass.h"
 #include "mlir/Conversion/ShapeToStandard/ShapeToStandard.h"
+#include "mlir/Conversion/ShardToMPI/ShardToMPI.h"
 #include "mlir/Conversion/TensorToLinalg/TensorToLinalgPass.h"
 #include "mlir/Conversion/TensorToSPIRV/TensorToSPIRVPass.h"
 #include "mlir/Conversion/TosaToArith/TosaToArith.h"

mlir/include/mlir/Conversion/Passes.h

@@ -903,13 +903,13 @@ def ConvertMemRefToSPIRVPass : Pass<"convert-memref-to-spirv"> {
 }
 
 //===----------------------------------------------------------------------===//
-// MeshToMPI
+// ShardToMPI
 //===----------------------------------------------------------------------===//
 
-def ConvertMeshToMPIPass : Pass<"convert-mesh-to-mpi"> {
-  let summary = "Convert Mesh dialect to MPI dialect.";
+def ConvertShardToMPIPass : Pass<"convert-shard-to-mpi"> {
+  let summary = "Convert Shard dialect to MPI dialect.";
   let description = [{
-    This pass converts communication operations from the Mesh dialect to the
+    This pass converts communication operations from the Shard dialect to the
     MPI dialect.
     If it finds the DLTI attribute "MPI:comm_world-rank" on the module it will
     use that integer value instead of calling MPI_Comm_rank. This allows

mlir/include/mlir/Conversion/Passes.td

@@ -1,23 +1,23 @@
-//===- MeshToMPI.h - Convert Mesh to MPI dialect ----------------*- C++ -*-===//
+//===- ShardToMPI.h - Convert Shard to MPI dialect --------------*- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
-#ifndef MLIR_CONVERSION_MESHTOMPI_MESHTOMPI_H
-#define MLIR_CONVERSION_MESHTOMPI_MESHTOMPI_H
+#ifndef MLIR_CONVERSION_SHARDTOMPI_SHARDTOMPI_H
+#define MLIR_CONVERSION_SHARDTOMPI_SHARDTOMPI_H
 
 #include "mlir/Pass/Pass.h"
 #include "mlir/Support/LLVM.h"
 
 namespace mlir {
 class Pass;
 
-#define GEN_PASS_DECL_CONVERTMESHTOMPIPASS
+#define GEN_PASS_DECL_CONVERTSHARDTOMPIPASS
 #include "mlir/Conversion/Passes.h.inc"
 
 } // namespace mlir
 
-#endif // MLIR_CONVERSION_MESHTOMPI_MESHTOMPI_H
+#endif // MLIR_CONVERSION_SHARDTOMPI_SHARDTOMPI_H

mlir/include/mlir/Conversion/MeshToMPI/MeshToMPI.h renamed to mlir/include/mlir/Conversion/ShardToMPI/ShardToMPI.h

@@ -19,7 +19,7 @@ add_subdirectory(Linalg)
 add_subdirectory(LLVMIR)
 add_subdirectory(Math)
 add_subdirectory(MemRef)
-add_subdirectory(Mesh)
+add_subdirectory(Shard)
 add_subdirectory(MLProgram)
 add_subdirectory(MPI)
 add_subdirectory(NVGPU)

mlir/include/mlir/Dialect/CMakeLists.txt

@@ -1,4 +1,4 @@
-//===- MeshShardingExtensions.h - -----------------------------------------===//
+//===- ShardingExtensions.h - -----------------------------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.

mlir/include/mlir/Dialect/Func/Extensions/MeshShardingExtensions.h renamed to mlir/include/mlir/Dialect/Func/Extensions/ShardingExtensions.h

@@ -1,20 +1,20 @@
-//===- MeshShardingInterfaceImpl.h ----------------------------------------===//
+//===- ShardingInterfaceImpl.h ----------------------------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
-#ifndef MLIR_DIALECT_LINALG_MESHSHARDINGINTERFACEIMPL_H
-#define MLIR_DIALECT_LINALG_MESHSHARDINGINTERFACEIMPL_H
+#ifndef MLIR_DIALECT_LINALG_SHARDSHARDINGINTERFACEIMPL_H
+#define MLIR_DIALECT_LINALG_SHARDSHARDINGINTERFACEIMPL_H
 
 namespace mlir {
 class DialectRegistry;
 
 namespace linalg {
-void registerMeshShardingInterfaceExternalModels(DialectRegistry &registry);
+void registerShardingInterfaceExternalModels(DialectRegistry &registry);
 } // namespace linalg
 } // namespace mlir
 
-#endif // MLIR_DIALECT_LINALG_MESHSHARDINGINTERFACEIMPL_H
+#endif // MLIR_DIALECT_LINALG_SHARDSHARDINGINTERFACEIMPL_H