[Mlir-commits] [mlir] [mlir][shard, mpi] Allowing 2d-grids and simplifying lowering shard.all_gather (PR #180243)

[Rolf Morel]

================
@@ -635,104 +634,124 @@ struct ConvertAllGatherOp : public CommOpPattern<AllGatherOp> {
     FailureOr<GridOp> gridOp = checkGrid(op, symbolTableCollection);
     if (failed(gridOp))
       return failure();
-    ImplicitLocOpBuilder iBuilder(op.getLoc(), rewriter);
-    Value input = getAsMemref(adaptor.getInput(), iBuilder);
+
+    ImplicitLocOpBuilder ib(op.getLoc(), rewriter);
+    Value input = getAsMemref(adaptor.getInput(), ib);
     MemRefType inType = cast<MemRefType>(input.getType());
-    if (!memref::isStaticShapeAndContiguousRowMajor(inType))
-      return op.emitError(
-          "Expected static shaped memref in contiguous row-major layout.");
     MemRefType outType = getMemrefType(cast<ShapedType>(op.getType()));
-    if (!memref::isStaticShapeAndContiguousRowMajor(outType))
-      return op.emitError(
-          "Expected static shaped memref in contiguous row-major layout.");
+    auto inputShape = inType.getShape();
+    auto outputShape = outType.getShape();
     int64_t gatherAxis = adaptor.getGatherAxisAttr().getInt();
-    auto ctx = op->getContext();
-
-    // Get the right communicator
-    Value comm = getComm(*gridOp, adaptor.getGridAxes(), iBuilder);
-
-    Value nRanks =
-        mpi::CommSizeOp::create(iBuilder, iBuilder.getI32Type(), comm)
-            .getSize();
-    nRanks =
-        arith::IndexCastOp::create(iBuilder, iBuilder.getIndexType(), nRanks);
-
-    Value tmpOutput, gatherDimSz;
-    if (gatherAxis == 0) {
-      tmpOutput = memref::AllocOp::create(iBuilder, outType);
-    } else {
-      // MPI's allgather always concatenates along the first dimension.
-      // Create a memref type for the output buffer with adjusted (expanded)
-      // shape.
-      SmallVector<int64_t> gatherShape(1, ShapedType::kDynamic);
-      llvm::append_range(gatherShape, outType.getShape());
-      gatherShape[gatherAxis + 1] = ShapedType::kDynamic;
-      MemRefType gatherType =
-          MemRefType::get(gatherShape, outType.getElementType());
-      gatherDimSz = arith::ConstantIndexOp::create(
-          iBuilder, outType.getDimSize(gatherAxis));
-      gatherDimSz = arith::DivSIOp::create(iBuilder, iBuilder.getIndexType(),
-                                           gatherDimSz, nRanks);
-      // Allocate output buffer
-      tmpOutput =
-          memref::AllocOp::create(iBuilder, gatherType, {nRanks, gatherDimSz});
+    int64_t inputDimOnAxis = inputShape[gatherAxis];
+    int64_t outputDimOnAxis = outputShape[gatherAxis];
+
+    for (size_t i = 0; i < outputShape.size(); ++i)
+      if (outputShape[i] != inputShape[i] && i != (size_t)gatherAxis)
+        return op.emitError(
+            "Result and input shapes must match along non-gather axes.");
+    if (inputDimOnAxis == 0)
+      return op.emitError("Input size along the gather axis must be non-zero.");
+    if (inputDimOnAxis == 1) {
+      assert(outputDimOnAxis == inputDimOnAxis);
+      rewriter.replaceOp(op, adaptor.getInput());
+      return success();
     }
+    if (outputDimOnAxis % inputDimOnAxis != 0)
+      return op.emitError("Result size along the gather axis must be an exact "
+                          "multiple of the input size along the gather axis.");
+
+    if (!memref::isStaticShapeAndContiguousRowMajor(inType) ||
+        !memref::isStaticShapeAndContiguousRowMajor(outType))
+      return op.emitError("Input/result must be statically shaped memrefs in "
+                          "contiguous row-major layout.");
+
+    // Get the right communicator.
+    Value comm = getComm(*gridOp, adaptor.getGridAxes(), ib);
+    Value nRanksV =
+        mpi::CommSizeOp::create(ib, ib.getI32Type(), comm).getSize();
+    nRanksV = arith::IndexCastOp::create(ib, ib.getIndexType(), nRanksV);
+    int64_t nRanks = outputDimOnAxis / inputDimOnAxis;
+    Value nRanksC = arith::ConstantIndexOp::create(ib, nRanks);
+    Value notError =
+        arith::CmpIOp::create(ib, arith::CmpIPredicate::eq, nRanksV, nRanksC);
+    cf::AssertOp::create(ib, notError,
+                         "Expected number of ranks in the communicator to "
+                         "match the output size along the gather axis divided "
+                         "by the input size along the gather axis.");
+
+    // mpi.allgather always concatenates along the first dimension, so
+    // get a output buffer of shape {nRanks, dim0, ...}.
+    SmallVector<int64_t> gatherShape;
+    gatherShape.emplace_back(nRanks);
+    gatherShape.append(inputShape.begin(), inputShape.end());
+    auto gatherType = MemRefType::get(gatherShape, outType.getElementType());
+    Value finalOutput = memref::AllocOp::create(ib, gatherType);
     // Create the MPI AllGather operation.
-    mpi::AllGatherOp::create(iBuilder, TypeRange(), input, tmpOutput, comm);
-
-    // If gather-axis!=0, copy from gathered buffer to output with the right
-    // layout.
-    Value finalOutput = tmpOutput;
-    if (gatherAxis != 0) {
-      int64_t nSrcDims = cast<ShapedType>(tmpOutput.getType()).getRank();
-      assert(nSrcDims == outType.getRank() + 1 &&
-             "Expected gathered type to have rank one more than output type.");
-
-      // Create affine map for copying from gathered buffer to output.
-      SmallVector<AffineExpr> dims;
-      dims.reserve(nSrcDims);
-      for (unsigned i = 0; i < nSrcDims; ++i)
-        dims.emplace_back(getAffineDimExpr(i, ctx));
-      AffineExpr s = getAffineSymbolExpr(0, ctx);
-      SmallVector<AffineExpr> results;
-      results.reserve(nSrcDims);
-      for (unsigned i = 0; i < nSrcDims - 1; ++i) {
-        if (i == gatherAxis)
-          results.emplace_back(dims[0] * s + dims[gatherAxis + 1]);
-        else
-          results.emplace_back(dims[i + 1]);
+    mpi::AllGatherOp::create(ib, TypeRange(), input, finalOutput, comm);
+
+    if (gatherAxis == 0) {
+      // If gather axis == 0, simply collapse the first 2 dims from {nRanks,
+      // dim0, ...} to {nRanks*dim0, ...}.
+      SmallVector<ReassociationIndices> reassociation;
+      reassociation.push_back({0, 1});
+      for (int64_t i = 2; i < (int64_t)gatherShape.size(); ++i) {
+        reassociation.push_back({i});
       }
----------------
rolfmorel wrote:

```suggestion
      int64_t numGatherDims = gatherShape.size();
      for (int64_t i = 2; i < numGatherDims; ++i)
        reassociation.push_back({i});
```
Nit: per the style guide.

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