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

[Rolf Morel]

================
@@ -645,94 +645,101 @@ struct ConvertAllGatherOp : public CommOpPattern<AllGatherOp> {
     if (!memref::isStaticShapeAndContiguousRowMajor(outType))
       return op.emitError(
           "Expected static shaped memref in contiguous row-major layout.");
-    int64_t gatherAxis = adaptor.getGatherAxisAttr().getInt();
-    auto ctx = op->getContext();
 
-    // Get the right communicator
-    Value comm = getComm(*gridOp, adaptor.getGridAxes(), iBuilder);
+    auto inputShape = cast<ShapedType>(adaptor.getInput().getType()).getShape();
+    auto outputShape = outType.getShape();
+    int64_t gatherAxis = adaptor.getGatherAxisAttr().getInt();
 
-    Value nRanks =
-        mpi::CommSizeOp::create(iBuilder, iBuilder.getI32Type(), comm)
-            .getSize();
-    nRanks =
-        arith::IndexCastOp::create(iBuilder, iBuilder.getIndexType(), nRanks);
+    // 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 = outputShape[gatherAxis] / inputShape[gatherAxis];
+    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.");
+    for (size_t i = 0; i < outputShape.size(); ++i)
+      assert(outputShape[i] == inputShape[i] || i == (size_t)gatherAxis);
+
+    // 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(ib, TypeRange(), input, finalOutput, comm);
 
-    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});
-    }
-    // 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]);
+      // 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});
       }
-      auto affineMap = AffineMap::get(nSrcDims, /*symbols=*/1, results, ctx);
-
-      finalOutput = memref::AllocOp::create(iBuilder, outType);
-
-      // Now build a loop nest to copy from gathered buffer to finalOutput
-      // It would be nicer to just use a memref.transpose/collapse_shape op but
-      // these currently only support simpler cases.
-      Value zero = arith::ConstantIndexOp::create(iBuilder, 0);
-      SmallVector<Value> lbs(nSrcDims, zero);
-      SmallVector<Value> ubs;
-      for (int64_t d = 0; d < nSrcDims; ++d)
-        ubs.emplace_back(memref::DimOp::create(iBuilder, tmpOutput, d));
-      SmallVector<int64_t> steps(nSrcDims, 1);
-      auto emitCopy = [&](OpBuilder &builder, Location loc, ValueRange ivs) {
-        Value v = memref::LoadOp::create(iBuilder, tmpOutput, ivs);
-        // set symbol value
-        SmallVector<Value> ivss(ivs.begin(), ivs.end());
-        ivss.emplace_back(gatherDimSz);
-        affine::AffineStoreOp::create(iBuilder, v, finalOutput, affineMap,
-                                      ivss);
-      };
-      affine::buildAffineLoopNest(iBuilder, op->getLoc(), lbs, ubs, steps,
-                                  emitCopy);
+      finalOutput = memref::CollapseShapeOp::create(ib, outType, finalOutput,
+                                                    reassociation);
 
-      memref::DeallocOp::create(iBuilder, tmpOutput);
+      // If the op's result is a tensor, cast it to a tensor.
+      if (isa<RankedTensorType>(op.getType()))
+        finalOutput = bufferization::ToTensorOp::create(ib, op.getType(),
+                                                        finalOutput, true);
+    } else {
+      // 1. Enter tensor-land.
+      auto inType =
+          RankedTensorType::get(gatherShape, outType.getElementType());
+      finalOutput =
+          bufferization::ToTensorOp::create(ib, inType, finalOutput, true);
+
+      // 2. Permute the output buffer from {nRanks, dim0, ..., gatherAxis, ...}
+      // to {dim0, ..., nRanks, dim1,...}.
+      SmallVector<int64_t> outShapePermuted, permutation;
----------------
rolfmorel wrote:

I think Copilot is on to something: could you avoid going back to tensors and just do the same transpose and collapse on memrefs?

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