[Mlir-commits] [mlir] [MLIR][Shard] Fold all_slice(all_gather(...)) pairs (PR #193906)

[llvmlistbot at llvm.org]

https://github.com/zackc6 updated https://github.com/llvm/llvm-project/pull/193906

>From fafb02cef993d6aa43afb8f0fdc993b4737317cf Mon Sep 17 00:00:00 2001
From: zack <zackchen666 at gmail.com>
Date: Fri, 24 Apr 2026 14:54:34 +0800
Subject: [PATCH] [MLIR][Shard] Fold all_gather/all_slice inverse pairs

Add a simplify pattern that replaces all_gather(all_slice(x)) with x when grid, grid axes, and gather/slice axis match, complementing the existing all_slice(all_gather(x)) fold. Extend shard simplify tests with positive and negative cases for both directions.
---
 .../lib/Dialect/Shard/Transforms/Simplify.cpp |  47 +++++++-
 mlir/test/Dialect/Shard/simplify.mlir         | 102 ++++++++++++++++++
 2 files changed, 148 insertions(+), 1 deletion(-)

diff --git a/mlir/lib/Dialect/Shard/Transforms/Simplify.cpp b/mlir/lib/Dialect/Shard/Transforms/Simplify.cpp
index 525ff007bc2f6..e948ecde4d7e8 100644
--- a/mlir/lib/Dialect/Shard/Transforms/Simplify.cpp
+++ b/mlir/lib/Dialect/Shard/Transforms/Simplify.cpp
@@ -131,6 +131,50 @@ struct AllReduceAllSliceSimplification : OpRewritePattern<AllSliceOp> {
   }
 };
 
+// Simplify AllSliceOp(AllGatherOp) -> input when both ops share the same grid,
+// grid_axes and axis. all_gather replicates grouped slices along gather_axis
+// and all_slice immediately picks the per-rank slice back out on the same axis.
+struct AllGatherAllSliceSimplification : OpRewritePattern<AllSliceOp> {
+  using OpRewritePattern::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(AllSliceOp sliceOp,
+                                PatternRewriter &rewriter) const override {
+    auto gatherOp = sliceOp.getInput().getDefiningOp<AllGatherOp>();
+    if (!gatherOp)
+      return failure();
+
+    if (gatherOp.getGrid() != sliceOp.getGrid() ||
+        gatherOp.getGridAxes() != sliceOp.getGridAxes() ||
+        gatherOp.getGatherAxis() != sliceOp.getSliceAxis())
+      return failure();
+
+    rewriter.replaceOp(sliceOp, gatherOp.getInput());
+    return success();
+  }
+};
+
+// Simplify AllGatherOp(AllSliceOp) -> input when both ops share the same grid,
+// grid_axes and axis. all_slice extracts each rank's shard along slice_axis and
+// all_gather on the same axis immediately stitches these shards back together.
+struct AllSliceAllGatherSimplification : OpRewritePattern<AllGatherOp> {
+  using OpRewritePattern::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(AllGatherOp gatherOp,
+                                PatternRewriter &rewriter) const override {
+    auto sliceOp = gatherOp.getInput().getDefiningOp<AllSliceOp>();
+    if (!sliceOp)
+      return failure();
+
+    if (sliceOp.getGrid() != gatherOp.getGrid() ||
+        sliceOp.getGridAxes() != gatherOp.getGridAxes() ||
+        sliceOp.getSliceAxis() != gatherOp.getGatherAxis())
+      return failure();
+
+    rewriter.replaceOp(gatherOp, sliceOp.getInput());
+    return success();
+  }
+};
+
 } // namespace
 
 void populateSimplifyPatterns(RewritePatternSet &patterns,
@@ -154,7 +198,8 @@ void populateSimplifyPatterns(RewritePatternSet &patterns,
   populateAllReduceEndomorphismSimplifyPatterns<arith::MaxUIOp>(
       patterns, ReductionKind::Max);
 
-  patterns.add<AllReduceAllSliceSimplification>(patterns.getContext());
+  patterns.add<AllReduceAllSliceSimplification, AllGatherAllSliceSimplification,
+               AllSliceAllGatherSimplification>(patterns.getContext());
 
   // TODO: add simplify patterns for all-gather and other collectives.
 
diff --git a/mlir/test/Dialect/Shard/simplify.mlir b/mlir/test/Dialect/Shard/simplify.mlir
index e5693a288fda6..181ccde98c505 100644
--- a/mlir/test/Dialect/Shard/simplify.mlir
+++ b/mlir/test/Dialect/Shard/simplify.mlir
@@ -1,3 +1,105 @@
+// RUN: mlir-opt %s -shard-simplify | FileCheck %s
+
+shard.grid @grid_ag(shape = 2x2)
+shard.grid @grid_ag_alt(shape = 2x2)
+
+// CHECK-LABEL: func.func @all_gather_all_slice_identity
+func.func @all_gather_all_slice_identity(
+    %arg0: tensor<4x4xf32>) -> tensor<4x4xf32> {
+  %0 = shard.all_gather %arg0 on @grid_ag grid_axes = [1] gather_axis = 1
+    : tensor<4x4xf32> -> tensor<4x8xf32>
+  %1 = shard.all_slice %0 on @grid_ag grid_axes = [1] slice_axis = 1
+    : tensor<4x8xf32> -> tensor<4x4xf32>
+  // CHECK-NOT: shard.all_gather
+  // CHECK-NOT: shard.all_slice
+  // CHECK: return %arg0 : tensor<4x4xf32>
+  return %1 : tensor<4x4xf32>
+}
+
+// CHECK-LABEL: func.func @all_gather_all_slice_different_axis
+func.func @all_gather_all_slice_different_axis(
+    %arg0: tensor<4x4xf32>) -> tensor<2x8xf32> {
+  %0 = shard.all_gather %arg0 on @grid_ag grid_axes = [1] gather_axis = 1
+    : tensor<4x4xf32> -> tensor<4x8xf32>
+  %1 = shard.all_slice %0 on @grid_ag grid_axes = [1] slice_axis = 0
+    : tensor<4x8xf32> -> tensor<2x8xf32>
+  // CHECK: shard.all_gather
+  // CHECK: shard.all_slice
+  return %1 : tensor<2x8xf32>
+}
+
+// CHECK-LABEL: func.func @all_gather_all_slice_different_grid_axes
+func.func @all_gather_all_slice_different_grid_axes(
+    %arg0: tensor<4x4xf32>) -> tensor<4x4xf32> {
+  %0 = shard.all_gather %arg0 on @grid_ag grid_axes = [0] gather_axis = 0
+    : tensor<4x4xf32> -> tensor<8x4xf32>
+  %1 = shard.all_slice %0 on @grid_ag grid_axes = [1] slice_axis = 0
+    : tensor<8x4xf32> -> tensor<4x4xf32>
+  // CHECK: shard.all_gather
+  // CHECK: shard.all_slice
+  return %1 : tensor<4x4xf32>
+}
+
+// CHECK-LABEL: func.func @all_gather_all_slice_different_grid
+func.func @all_gather_all_slice_different_grid(
+    %arg0: tensor<4x4xf32>) -> tensor<4x4xf32> {
+  %0 = shard.all_gather %arg0 on @grid_ag grid_axes = [1] gather_axis = 1
+    : tensor<4x4xf32> -> tensor<4x8xf32>
+  %1 = shard.all_slice %0 on @grid_ag_alt grid_axes = [1] slice_axis = 1
+    : tensor<4x8xf32> -> tensor<4x4xf32>
+  // CHECK: shard.all_gather
+  // CHECK: shard.all_slice
+  return %1 : tensor<4x4xf32>
+}
+
+// CHECK-LABEL: func.func @all_slice_all_gather_identity
+func.func @all_slice_all_gather_identity(
+    %arg0: tensor<4x4xf32>) -> tensor<4x4xf32> {
+  %0 = shard.all_slice %arg0 on @grid_ag grid_axes = [1] slice_axis = 1
+    : tensor<4x4xf32> -> tensor<4x2xf32>
+  %1 = shard.all_gather %0 on @grid_ag grid_axes = [1] gather_axis = 1
+    : tensor<4x2xf32> -> tensor<4x4xf32>
+  // CHECK-NOT: shard.all_slice
+  // CHECK-NOT: shard.all_gather
+  // CHECK: return %arg0 : tensor<4x4xf32>
+  return %1 : tensor<4x4xf32>
+}
+
+// CHECK-LABEL: func.func @all_slice_all_gather_different_axis
+func.func @all_slice_all_gather_different_axis(
+    %arg0: tensor<4x4xf32>) -> tensor<8x2xf32> {
+  %0 = shard.all_slice %arg0 on @grid_ag grid_axes = [1] slice_axis = 1
+    : tensor<4x4xf32> -> tensor<4x2xf32>
+  %1 = shard.all_gather %0 on @grid_ag grid_axes = [1] gather_axis = 0
+    : tensor<4x2xf32> -> tensor<8x2xf32>
+  // CHECK: shard.all_slice
+  // CHECK: shard.all_gather
+  return %1 : tensor<8x2xf32>
+}
+
+// CHECK-LABEL: func.func @all_slice_all_gather_different_grid
+func.func @all_slice_all_gather_different_grid(
+    %arg0: tensor<4x4xf32>) -> tensor<4x4xf32> {
+  %0 = shard.all_slice %arg0 on @grid_ag grid_axes = [1] slice_axis = 1
+    : tensor<4x4xf32> -> tensor<4x2xf32>
+  %1 = shard.all_gather %0 on @grid_ag_alt grid_axes = [1] gather_axis = 1
+    : tensor<4x2xf32> -> tensor<4x4xf32>
+  // CHECK: shard.all_slice
+  // CHECK: shard.all_gather
+  return %1 : tensor<4x4xf32>
+}
+
+// CHECK-LABEL: func.func @all_slice_all_gather_different_grid_axes
+func.func @all_slice_all_gather_different_grid_axes(
+    %arg0: tensor<4x4xf32>) -> tensor<4x4xf32> {
+  %0 = shard.all_slice %arg0 on @grid_ag grid_axes = [0] slice_axis = 0
+    : tensor<4x4xf32> -> tensor<2x4xf32>
+  %1 = shard.all_gather %0 on @grid_ag grid_axes = [1] gather_axis = 0
+    : tensor<2x4xf32> -> tensor<4x4xf32>
+  // CHECK: shard.all_slice
+  // CHECK: shard.all_gather
+  return %1 : tensor<4x4xf32>
+}
 // RUN: mlir-opt -shard-simplify %s | FileCheck %s
 
 shard.grid @grid0(shape = 4x2)