TheRock/compiler/amd-llvm/mlir/lib/Conversion/ShardToMPI/ShardToMPI.cpp

//===- ShardToMPI.cpp - Shard to MPI  dialect conversion -----------------===//
//
// 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
//
//===----------------------------------------------------------------------===//
//
// This file implements a translation of Shard communication ops to MPI ops.
//
//===----------------------------------------------------------------------===//

#include "mlir/Conversion/ShardToMPI/ShardToMPI.h"
#include "mlir/Dialect/Shard/Transforms/Transforms.h"

#include "mlir/Dialect/Affine/IR/AffineOps.h"
#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/Bufferization/IR/Bufferization.h"
#include "mlir/Dialect/ControlFlow/IR/ControlFlowOps.h"
#include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/Dialect/Func/Transforms/FuncConversions.h"
#include "mlir/Dialect/Linalg/IR/Linalg.h"
#include "mlir/Dialect/MPI/IR/MPI.h"
#include "mlir/Dialect/MemRef/IR/MemRef.h"
#include "mlir/Dialect/MemRef/Utils/MemRefUtils.h"
#include "mlir/Dialect/SCF/IR/SCF.h"
#include "mlir/Dialect/Shard/IR/ShardDialect.h"
#include "mlir/Dialect/Shard/IR/ShardOps.h"
#include "mlir/Dialect/Shard/Transforms/Simplify.h"
#include "mlir/Dialect/Shard/Transforms/Transforms.h"
#include "mlir/Dialect/Tensor/IR/Tensor.h"
#include "mlir/Dialect/Utils/StaticValueUtils.h"
#include "mlir/IR/AffineMap.h"
#include "mlir/IR/Builders.h"
#include "mlir/IR/BuiltinAttributes.h"
#include "mlir/IR/BuiltinTypes.h"
#include "mlir/IR/PatternMatch.h"
#include "mlir/IR/SymbolTable.h"
#include "mlir/Transforms/DialectConversion.h"
#include "mlir/Transforms/GreedyPatternRewriteDriver.h"

#define DEBUG_TYPE "shard-to-mpi"

namespace mlir {
#define GEN_PASS_DEF_CONVERTSHARDTOMPIPASS
#include "mlir/Conversion/Passes.h.inc"
} // namespace mlir

using namespace mlir;
using namespace shard;

namespace {
/// Converts a vector of OpFoldResults (ints) into vector of Values of the
/// provided type.
static SmallVector<Value> getMixedAsValues(OpBuilder b, const Location &loc,
                                           llvm::ArrayRef<int64_t> statics,
                                           ValueRange dynamics,
                                           Type type = Type()) {
  SmallVector<Value> values;
  auto dyn = dynamics.begin();
  Type i64 = b.getI64Type();
  if (!type)
    type = i64;
  assert((i64 == type || b.getIndexType() == type) &&
         "expected an i64 or an intex type");
  for (auto s : statics) {
    if (s == ShapedType::kDynamic) {
      values.emplace_back(*(dyn++));
    } else {
      TypedAttr val = type == i64 ? b.getI64IntegerAttr(s) : b.getIndexAttr(s);
      values.emplace_back(arith::ConstantOp::create(b, loc, type, val));
    }
  }
  return values;
}

/// Create operations converting a linear index to a multi-dimensional index.
[[maybe_unused]] static SmallVector<Value>
linearToMultiIndex(Location loc, OpBuilder b, Value linearIndex,
                   ValueRange dimensions) {
  int n = dimensions.size();
  SmallVector<Value> multiIndex(n);

  for (int i = n - 1; i >= 0; --i) {
    multiIndex[i] = arith::RemSIOp::create(b, loc, linearIndex, dimensions[i]);
    if (i > 0)
      linearIndex = arith::DivSIOp::create(b, loc, linearIndex, dimensions[i]);
  }

  return multiIndex;
}

/// Create operations converting a multi-dimensional index to a linear index.
Value multiToLinearIndex(Location loc, OpBuilder b, ValueRange multiIndex,
                         ValueRange dimensions) {

  Value linearIndex = arith::ConstantIndexOp::create(b, loc, 0);
  Value stride = arith::ConstantIndexOp::create(b, loc, 1);

  for (int i = multiIndex.size() - 1; i >= 0; --i) {
    Value off = arith::MulIOp::create(b, loc, multiIndex[i], stride);
    linearIndex = arith::AddIOp::create(b, loc, linearIndex, off);
    stride = arith::MulIOp::create(b, loc, stride, dimensions[i]);
  }

  return linearIndex;
}

/// Replace GetShardingOp with related/dependent ShardingOp.
struct ConvertGetShardingOp : public OpConversionPattern<GetShardingOp> {
  using OpConversionPattern::OpConversionPattern;

  LogicalResult
  matchAndRewrite(GetShardingOp op, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    auto shardOp = adaptor.getSource().getDefiningOp<ShardOp>();
    if (!shardOp)
      return failure();
    auto shardingOp = shardOp.getSharding().getDefiningOp<ShardingOp>();
    if (!shardingOp)
      return failure();

    rewriter.replaceOp(op, shardingOp.getResult());
    return success();
  }
};

/// Convert a sharding op to a tuple of tensors of its components
///   (SplitAxes, HaloSizes, ShardedDimsOffsets)
/// as defined by type converter.
struct ConvertShardingOp : public OpConversionPattern<ShardingOp> {
  using OpConversionPattern::OpConversionPattern;

  LogicalResult
  matchAndRewrite(ShardingOp op, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    auto splitAxes = op.getSplitAxes().getAxes();
    int64_t maxNAxes = 0;
    for (auto axes : splitAxes)
      maxNAxes = std::max<int64_t>(maxNAxes, axes.size());

    // To hold the split axes, create empty 2d tensor with shape
    // {splitAxes.size(), max-size-of-split-groups}.
    // Set trailing elements for smaller split-groups to -1.
    Location loc = op.getLoc();
    auto i16 = rewriter.getI16Type();
    auto i64 = rewriter.getI64Type();
    std::array<int64_t, 2> shape = {static_cast<int64_t>(splitAxes.size()),
                                    maxNAxes};
    Value resSplitAxes = tensor::EmptyOp::create(rewriter, loc, shape, i16);
    auto attr = IntegerAttr::get(i16, -1);
    Value fillValue = arith::ConstantOp::create(rewriter, loc, i16, attr);
    resSplitAxes =
        linalg::FillOp::create(rewriter, loc, fillValue, resSplitAxes)
            .getResult(0);

    // explicitly write values into tensor row by row
    std::array<int64_t, 2> strides = {1, 1};
    int64_t nSplits = 0;
    ValueRange empty = {};
    for (auto [i, axes] : llvm::enumerate(splitAxes)) {
      int64_t size = axes.size();
      if (size > 0)
        ++nSplits;
      std::array<int64_t, 2> offs = {(int64_t)i, 0};
      std::array<int64_t, 2> sizes = {1, size};
      auto tensorType = RankedTensorType::get({size}, i16);
      auto attrs = DenseIntElementsAttr::get(tensorType, axes.asArrayRef());
      auto vals = arith::ConstantOp::create(rewriter, loc, tensorType, attrs);
      resSplitAxes = tensor::InsertSliceOp::create(rewriter, loc, vals,
                                                   resSplitAxes, empty, empty,
                                                   empty, offs, sizes, strides);
    }

    // To hold halos sizes, create 2d Tensor with shape {nSplits, 2}.
    // Store the halo sizes in the tensor.
    SmallVector<Value> haloSizes =
        getMixedAsValues(rewriter, loc, adaptor.getStaticHaloSizes(),
                         adaptor.getDynamicHaloSizes());
    auto type = RankedTensorType::get({nSplits, 2}, i64);
    Value resHaloSizes =
        haloSizes.empty()
            ? tensor::EmptyOp::create(rewriter, loc,
                                      std::array<int64_t, 2>{0, 0}, i64)
                  .getResult()
            : tensor::FromElementsOp::create(rewriter, loc, type, haloSizes)
                  .getResult();

    // To hold sharded dims offsets, create Tensor with shape {nSplits,
    // maxSplitSize+1}. Store the offsets in the tensor but set trailing
    // elements for smaller split-groups to -1. Computing the max size of the
    // split groups needs using collectiveProcessGroupSize (which needs the
    // GridOp)
    Value resOffsets;
    if (adaptor.getStaticShardedDimsOffsets().empty()) {
      resOffsets = tensor::EmptyOp::create(rewriter, loc,
                                           std::array<int64_t, 2>{0, 0}, i64);
    } else {
      SymbolTableCollection symbolTableCollection;
      auto gridOp = getGrid(op, symbolTableCollection);
      int64_t maxSplitSize = 0;
      for (auto axes : splitAxes) {
        int64_t splitSize =
            collectiveProcessGroupSize(axes.asArrayRef(), gridOp.getShape());
        assert(splitSize != ShapedType::kDynamic);
        maxSplitSize = std::max<int64_t>(maxSplitSize, splitSize);
      }
      assert(maxSplitSize);
      ++maxSplitSize; // add one for the total size

      resOffsets = tensor::EmptyOp::create(
          rewriter, loc, std::array<int64_t, 2>{nSplits, maxSplitSize}, i64);
      Value zero = arith::ConstantOp::create(
          rewriter, loc, i64, rewriter.getI64IntegerAttr(ShapedType::kDynamic));
      resOffsets =
          linalg::FillOp::create(rewriter, loc, zero, resOffsets).getResult(0);
      SmallVector<Value> offsets =
          getMixedAsValues(rewriter, loc, adaptor.getStaticShardedDimsOffsets(),
                           adaptor.getDynamicShardedDimsOffsets());
      int64_t curr = 0;
      for (auto [i, axes] : llvm::enumerate(splitAxes)) {
        int64_t splitSize =
            collectiveProcessGroupSize(axes.asArrayRef(), gridOp.getShape());
        assert(splitSize != ShapedType::kDynamic && splitSize < maxSplitSize);
        ++splitSize; // add one for the total size
        ArrayRef<Value> values(&offsets[curr], splitSize);
        Value vals = tensor::FromElementsOp::create(rewriter, loc, values);
        std::array<int64_t, 2> offs = {static_cast<int64_t>(i), 0};
        std::array<int64_t, 2> sizes = {1, splitSize};
        resOffsets = tensor::InsertSliceOp::create(rewriter, loc, vals,
                                                   resOffsets, empty, empty,
                                                   empty, offs, sizes, strides);
        curr += splitSize;
      }
    }

    // return a tuple of tensors as defined by type converter
    SmallVector<Type> resTypes;
    if (failed(getTypeConverter()->convertType(op.getResult().getType(),
                                               resTypes)))
      return failure();

    resSplitAxes =
        tensor::CastOp::create(rewriter, loc, resTypes[0], resSplitAxes);
    resHaloSizes =
        tensor::CastOp::create(rewriter, loc, resTypes[1], resHaloSizes);
    resOffsets = tensor::CastOp::create(rewriter, loc, resTypes[2], resOffsets);

    rewriter.replaceOpWithNewOp<UnrealizedConversionCastOp>(
        op, TupleType::get(op.getContext(), resTypes),
        ValueRange{resSplitAxes, resHaloSizes, resOffsets});

    return success();
  }
};

class ConvertProcessLinearIndexOp
    : public OpConversionPattern<ProcessLinearIndexOp> {

public:
  using OpConversionPattern::OpConversionPattern;

  LogicalResult
  matchAndRewrite(ProcessLinearIndexOp op, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    // Create mpi::CommRankOp
    Location loc = op.getLoc();
    auto *ctx = op.getContext();
    Value commWorld =
        mpi::CommWorldOp::create(rewriter, loc, mpi::CommType::get(ctx));
    auto rank = mpi::CommRankOp::create(
                    rewriter, loc,
                    TypeRange{mpi::RetvalType::get(ctx), rewriter.getI32Type()},
                    commWorld)
                    .getRank();
    rewriter.replaceOpWithNewOp<arith::IndexCastOp>(op, rewriter.getIndexType(),
                                                    rank);
    return success();
  }
};

struct ConvertNeighborsLinearIndicesOp
    : public OpConversionPattern<NeighborsLinearIndicesOp> {
  using OpConversionPattern::OpConversionPattern;

  LogicalResult
  matchAndRewrite(NeighborsLinearIndicesOp op, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {

    // Computes the neighbors indices along a split axis by simply
    // adding/subtracting 1 to the current index in that dimension.
    // Assigns -1 if neighbor is out of bounds.

    auto axes = adaptor.getSplitAxes();
    // For now only single axis sharding is supported
    if (axes.size() != 1)
      return failure();

    Location loc = op.getLoc();
    SymbolTableCollection symbolTableCollection;
    auto gridOp = getGrid(op, symbolTableCollection);
    auto mIdx = adaptor.getDevice();
    auto orgIdx = mIdx[axes[0]];
    SmallVector<Value> dims;
    llvm::transform(
        gridOp.getShape(), std::back_inserter(dims), [&](int64_t i) {
          return arith::ConstantIndexOp::create(rewriter, loc, i).getResult();
        });
    Value dimSz = dims[axes[0]];
    Value one = arith::ConstantIndexOp::create(rewriter, loc, 1);
    Value minus1 = arith::ConstantIndexOp::create(rewriter, loc, -1);
    Value atBorder =
        arith::CmpIOp::create(rewriter, loc, arith::CmpIPredicate::sle, orgIdx,
                              arith::ConstantIndexOp::create(rewriter, loc, 0));
    auto down = scf::IfOp::create(
        rewriter, loc, atBorder,
        [&](OpBuilder &builder, Location loc) {
          scf::YieldOp::create(builder, loc, minus1);
        },
        [&](OpBuilder &builder, Location loc) {
          SmallVector<Value> tmp = mIdx;
          tmp[axes[0]] =
              arith::SubIOp::create(rewriter, op.getLoc(), orgIdx, one)
                  .getResult();
          scf::YieldOp::create(builder, loc,
                               multiToLinearIndex(loc, rewriter, tmp, dims));
        });
    atBorder = arith::CmpIOp::create(
        rewriter, loc, arith::CmpIPredicate::sge, orgIdx,
        arith::SubIOp::create(rewriter, loc, dimSz, one).getResult());
    auto up = scf::IfOp::create(
        rewriter, loc, atBorder,
        [&](OpBuilder &builder, Location loc) {
          scf::YieldOp::create(builder, loc, minus1);
        },
        [&](OpBuilder &builder, Location loc) {
          SmallVector<Value> tmp = mIdx;
          tmp[axes[0]] =
              arith::AddIOp::create(rewriter, op.getLoc(), orgIdx, one);
          scf::YieldOp::create(builder, loc,
                               multiToLinearIndex(loc, rewriter, tmp, dims));
        });
    rewriter.replaceOp(op, ValueRange{down.getResult(0), up.getResult(0)});
    return success();
  }
};

struct ConvertShardShapeOp : public OpConversionPattern<ShardShapeOp> {
  using OpConversionPattern::OpConversionPattern;

  LogicalResult
  matchAndRewrite(ShardShapeOp op, OneToNOpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    auto sharding = op.getSharding().getDefiningOp<ShardingOp>();
    if (!sharding) {
      return op->emitError()
             << "Expected ShardingOp as defining op for sharding"
             << " but found " << adaptor.getSharding()[0].getDefiningOp();
    }

    // Compute the sharded shape by applying the sharding to the input shape.
    // If shardedDimsOffsets is not defined in the sharding, the shard shape is
    // computed by dividing the dimension size by the number of shards in that
    // dimension (which is given by the size of the grid axes provided in
    // split-axes). Odd elements get distributed to trailing shards. If a
    // shardedDimsOffsets is provided, the shard shape is computed by
    // subtracting the offset of the current shard from the offset of the next
    // shard.

    Location loc = op.getLoc();
    Type index = rewriter.getIndexType();

    // This is a 1:N conversion because the sharding op is a 1:3 conversion.
    // The operands in the adaptor are a vector<ValeRange>. For dims and device
    // we have a 1:1 conversion.
    // For simpler access fill a vector with the dynamic dims.
    SmallVector<Value> dynDims, dynDevice;
    for (auto dim : adaptor.getDimsDynamic()) {
      // type conversion should be 1:1 for ints
      dynDims.emplace_back(llvm::getSingleElement(dim));
    }
    // same for device
    for (auto device : adaptor.getDeviceDynamic()) {
      dynDevice.emplace_back(llvm::getSingleElement(device));
    }

    // To keep the code simple, convert dims/device to values when they are
    // attributes. Count on canonicalization to fold static values.
    SmallVector<Value> shape =
        getMixedAsValues(rewriter, loc, op.getDims(), dynDims, index);
    SmallVector<Value> multiIdx =
        getMixedAsValues(rewriter, loc, adaptor.getDevice(), dynDevice, index);

    // Get the GridOp, the grid shape is needed to compute the sharded shape.
    SymbolTableCollection symbolTableCollection;
    auto gridOp = getGrid(sharding, symbolTableCollection);
    // For now we only support static grid shapes
    if (ShapedType::isDynamicShape(gridOp.getShape()))
      return failure();

    auto splitAxes = sharding.getSplitAxes().getAxes();
    // shardedDimsOffsets are optional and might be Values (not attributes).
    // Also, the shardId might be dynamic which means the position in the
    // shardedDimsOffsets is not statically known. Create a tensor of the
    // shardedDimsOffsets and later extract the offsets for computing the
    // local shard-size.
    Value shardedDimsOffs;
    {
      SmallVector<Value> tmp = getMixedAsValues(
          rewriter, loc, sharding.getStaticShardedDimsOffsets(),
          sharding.getDynamicShardedDimsOffsets(), index);
      if (!tmp.empty())
        shardedDimsOffs = tensor::FromElementsOp::create(
            rewriter, loc, RankedTensorType::get({(int64_t)tmp.size()}, index),
            tmp);
    }

    // With static grid shape the sizes of the split axes are known.
    // Hence the start/pos for each split axes in shardDimsOffsets can be
    // computed statically.
    int64_t pos = 0;
    SmallVector<Value> shardShape;
    Value zero =
        arith::ConstantOp::create(rewriter, loc, rewriter.getZeroAttr(index));
    Value one =
        arith::ConstantOp::create(rewriter, loc, rewriter.getOneAttr(index));

    // Iterate over the dimensions of the tensor shape, get their split Axes,
    // and compute the sharded shape.
    for (auto [i, dim] : llvm::enumerate(shape)) {
      // Trailing dimensions might not be annotated.
      if (i < splitAxes.size() && !splitAxes[i].empty()) {
        auto axes = splitAxes[i];
        // The current dimension might not be sharded.
        // Create a value from the static position in shardDimsOffsets.
        Value posVal = arith::ConstantOp::create(rewriter, loc,
                                                 rewriter.getIndexAttr(pos));
        // Get the index of the local shard in the grid axis.
        Value idx = multiIdx[axes[0]];
        auto numShards =
            collectiveProcessGroupSize(axes.asArrayRef(), gridOp.getShape());
        if (shardedDimsOffs) {
          // If sharded dims offsets are provided, use them to compute the
          // sharded shape.
          if (axes.size() > 1) {
            return op->emitError() << "Only single axis sharding is "
                                   << "supported for each dimension.";
          }
          idx = arith::AddIOp::create(rewriter, loc, posVal, idx);
          // Compute size = shardedDimsOffs[idx+1] - shardedDimsOffs[idx].
          Value off =
              tensor::ExtractOp::create(rewriter, loc, shardedDimsOffs, idx);
          idx = arith::AddIOp::create(rewriter, loc, idx, one);
          Value nextOff =
              tensor::ExtractOp::create(rewriter, loc, shardedDimsOffs, idx);
          Value sz = arith::SubIOp::create(rewriter, loc, nextOff, off);
          shardShape.emplace_back(sz);
        } else {
          Value numShardsVal = arith::ConstantOp::create(
              rewriter, loc, rewriter.getIndexAttr(numShards));
          // Compute shard dim size by distributing odd elements to trailing
          // shards:
          // sz = dim / numShards
          //      + (idx >= (numShards - (dim % numShards)) ? 1 : 0)
          Value sz = arith::DivSIOp::create(rewriter, loc, dim, numShardsVal);
          Value sz1 = arith::RemSIOp::create(rewriter, loc, dim, numShardsVal);
          sz1 = arith::SubIOp::create(rewriter, loc, numShardsVal, sz1);
          auto cond = arith::CmpIOp::create(
              rewriter, loc, arith::CmpIPredicate::sge, idx, sz1);
          Value odd = arith::SelectOp::create(rewriter, loc, cond, one, zero);
          sz = arith::AddIOp::create(rewriter, loc, sz, odd);
          shardShape.emplace_back(sz);
        }
        pos += numShards + 1; // add one for the total size.
      } // else no sharding if split axis is empty or no split axis
      // If no size was added -> no sharding in this dimension.
      if (shardShape.size() <= i)
        shardShape.emplace_back(dim);
    }
    assert(shardShape.size() == shape.size());
    rewriter.replaceOp(op, shardShape);
    return success();
  }
};

static mpi::MPI_ReductionOpEnumAttr getMPIReductionOp(ReductionKindAttr kind) {
  auto *ctx = kind.getContext();
  auto getReductionOp = [ctx](mpi::MPI_ReductionOpEnum redOp) {
    return mpi::MPI_ReductionOpEnumAttr::get(ctx, redOp);
  };

  switch (kind.getValue()) {
  case ReductionKind::Sum:
    return getReductionOp(mpi::MPI_ReductionOpEnum::MPI_SUM);
  case ReductionKind::Product:
    return getReductionOp(mpi::MPI_ReductionOpEnum::MPI_PROD);
  case ReductionKind::Min:
    return getReductionOp(mpi::MPI_ReductionOpEnum::MPI_MIN);
  case ReductionKind::Max:
    return getReductionOp(mpi::MPI_ReductionOpEnum::MPI_MAX);
  case ReductionKind::BitwiseAnd:
    return getReductionOp(mpi::MPI_ReductionOpEnum::MPI_BAND);
  case ReductionKind::BitwiseOr:
    return getReductionOp(mpi::MPI_ReductionOpEnum::MPI_BOR);
  case ReductionKind::BitwiseXor:
    return getReductionOp(mpi::MPI_ReductionOpEnum::MPI_BXOR);
  default:
    llvm_unreachable("Unknown/unsupported reduction kind");
  }
}

template <typename CommOp>
struct CommOpPattern : public OpConversionPattern<CommOp> {
  using OpConversionPattern<CommOp>::OpConversionPattern;

  MemRefType getMemrefType(ShapedType tensorType) const {
    return MemRefType::get(tensorType.getShape(), tensorType.getElementType());
  }

  Value getAsMemref(Value input, ImplicitLocOpBuilder &iBuilder,
                    bool readOnly) const {
    auto itype = input.getType();
    // If the source is a tensor, materialize a memref for it.
    if (isa<RankedTensorType>(itype)) {
      auto memrefType = getMemrefType(cast<ShapedType>(itype));
      input = bufferization::ToBufferOp::create(iBuilder, memrefType, input,
                                                readOnly);
    } else {
      assert(isa<MemRefType>(itype) &&
             "expected input to be of MemRefType or TensorType");
    }
    return input;
  }

  FailureOr<GridOp> checkGrid(CommOp op,
                              SymbolTableCollection &symbolTableCollection,
                              bool allowDynamic = false) const {
    GridOp gridOp = getGrid(op, symbolTableCollection);
    if (!gridOp)
      return op->emitError() << "Missing grid symbol.";
    if (!allowDynamic && ShapedType::isDynamicShape(gridOp.getShape()))
      return op->emitError() << "Dynamic grid shape not supported.";
    return gridOp;
  }

  // Get an MPI_Comm_split for a given grid and axes.
  // The color is the linear index of the process in the grid along the
  // non-'grid-axes'. The key is the linear index of the process in the grid
  // along the grid-axes.
  Value getComm(GridOp &gridOp, ::llvm::ArrayRef<int16_t> gridAxes,
                ImplicitLocOpBuilder &iBuilder) const {
    size_t gridDims = gridOp.getShape().size();
    auto commType = mpi::CommType::get(gridOp->getContext());
    Value commWorld = mpi::CommWorldOp::create(iBuilder, commType);

    if (gridAxes.empty() || gridAxes.size() >= gridDims) {
      return commWorld;
    }

    SmallVector<GridAxis> otherAxes;
    for (GridAxis i = 0; i < static_cast<GridAxis>(gridDims); ++i) {
      if (!llvm::is_contained(gridAxes, i))
        otherAxes.emplace_back(i);
    }

    SmallVector<Type> indexResultTypes(otherAxes.size(),
                                       iBuilder.getIndexType());

    Value color =
        createProcessLinearIndex(iBuilder, gridOp.getSymName(), otherAxes);
    color = arith::IndexCastOp::create(iBuilder, iBuilder.getI32Type(), color);

    Value key =
        createProcessLinearIndex(iBuilder, gridOp.getSymName(), gridAxes);
    key = arith::IndexCastOp::create(iBuilder, iBuilder.getI32Type(), key);

    // Finally split the communicator
    return mpi::CommSplitOp::create(iBuilder, commType, commWorld, color, key)
        .getNewcomm();
  }
};

struct ConvertAllReduceOp : public CommOpPattern<AllReduceOp> {
  using CommOpPattern::CommOpPattern;

  LogicalResult
  matchAndRewrite(AllReduceOp op, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    SymbolTableCollection symbolTableCollection;
    FailureOr<GridOp> gridOp = checkGrid(op, symbolTableCollection);
    if (failed(gridOp))
      return failure();
    ImplicitLocOpBuilder iBuilder(op.getLoc(), rewriter);
    Value input = getAsMemref(adaptor.getInput(), iBuilder, true);
    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.");

    // Allocate buffer and copy input to buffer.
    Value buffer = memref::AllocOp::create(iBuilder, outType);
    linalg::CopyOp::create(iBuilder, input, buffer);
    // Get the right communicator
    Value comm = getComm(*gridOp, adaptor.getGridAxes(), iBuilder);
    // Create the MPI AllReduce operation.
    mpi::AllReduceOp::create(iBuilder, TypeRange(), buffer, buffer,
                             getMPIReductionOp(adaptor.getReductionAttr()),
                             comm);

    // If the destination is a tensor, cast it to a tensor
    if (isa<RankedTensorType>(op.getType()))
      buffer = bufferization::ToTensorOp::create(iBuilder, op.getType(), buffer,
                                                 true);
    rewriter.replaceOp(op, buffer);
    return success();
  }
};

struct ConvertReduceScatterOp : public CommOpPattern<ReduceScatterOp> {
  using CommOpPattern::CommOpPattern;

  // shard.reduce_scatter reduces and then scatters along a specified
  // scatter-dim. mpi.reduce_scatter_block always scatters along the first
  // dimension. Hence, if scatter-dim != 0, we need to rearrange the input
  // data by expanding the scatter-dim into {nRanks, output_scatter_dim} and
  // transposing nRanks to the first dimension.

  LogicalResult
  matchAndRewrite(ReduceScatterOp op, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    auto gridAxes = adaptor.getGridAxes();
    int64_t scatterDim = adaptor.getScatterDimAttr().getInt();

    SymbolTableCollection symbolTableCollection;
    FailureOr<GridOp> gridOp = checkGrid(op, symbolTableCollection);
    if (failed(gridOp))
      return failure();

    ImplicitLocOpBuilder ib(op.getLoc(), rewriter);
    Value rawInput = adaptor.getInput();
    auto inShapedType = cast<ShapedType>(rawInput.getType());
    MemRefType outType = getMemrefType(cast<ShapedType>(op.getType()));
    auto elemType = outType.getElementType();
    auto inputShape = inShapedType.getShape();
    auto outputShape = outType.getShape();
    int64_t inputDimOnAxis = inputShape[scatterDim];
    int64_t outputDimOnAxis = outputShape[scatterDim];

    for (size_t i = 0; i < outputShape.size(); ++i)
      if (outputShape[i] != inputShape[i] &&
          i != static_cast<size_t>(scatterDim))
        return op.emitError(
            "Result and input shapes must match along non-scatter axes.");
    if (outputDimOnAxis == 0)
      return op.emitError(
          "Output size along the scatter axis must be non-zero.");
    if (inputDimOnAxis % outputDimOnAxis != 0)
      return op.emitError(
          "Input size along the scatter axis must be an exact "
          "multiple of the output size along the scatter axis.");

    if (!memref::isStaticShapeAndContiguousRowMajor(outType))
      return op.emitError("Result must be a statically shaped memref in "
                          "contiguous row-major layout.");

    int64_t nRanks = inputDimOnAxis / outputDimOnAxis;

    // Verify that nRanks matches the number of devices along the grid axes.
    int64_t gridGroupSize =
        collectiveProcessGroupSize(gridAxes, gridOp->getShape());
    if (nRanks != gridGroupSize)
      return op.emitError()
             << "Expected the scatter factor (" << nRanks
             << ") to match the number of devices along grid_axes ("
             << gridGroupSize << ").";

    // Get the right communicator.
    Value comm = getComm(*gridOp, gridAxes, ib);

    Value mpiInput;
    if (scatterDim == 0) {
      // scatter_dim == 0 maps directly to MPI_Reduce_scatter_block.
      // Input must be contiguous for MPI.
      Value input = getAsMemref(rawInput, ib, true);
      MemRefType inType = cast<MemRefType>(input.getType());
      if (!memref::isStaticShapeAndContiguousRowMajor(inType))
        return op.emitError("Input must be a statically shaped memref in "
                            "contiguous row-major layout.");
      mpiInput = input;
    } else {
      // For scatter_dim != 0 we rearrange the input so the scatter factor
      // becomes the first dimension.
      //
      // 1. Get a tensor representation of the input (avoid memref->tensor
      //    round-trip if the input is already a tensor).
      Value tensorInput = rawInput;
      if (!isa<RankedTensorType>(rawInput.getType())) {
        auto inTensorType = RankedTensorType::get(inputShape, elemType);
        tensorInput =
            bufferization::ToTensorOp::create(ib, inTensorType, rawInput, true);
      }

      // 2. Expand the scatter dim from {d0, ..., d_sd, ..., dN} to
      //    {d0, ..., nRanks, o_sd, ..., dN}.
      SmallVector<int64_t> expandedShape;
      SmallVector<ReassociationIndices> expandReassociation;
      int64_t expandedIdx = 0;
      for (int64_t i = 0; i < static_cast<int64_t>(inputShape.size()); ++i) {
        if (i == scatterDim) {
          expandedShape.push_back(nRanks);
          expandedShape.push_back(outputDimOnAxis);
          expandReassociation.push_back({expandedIdx, expandedIdx + 1});
          expandedIdx += 2;
        } else {
          expandedShape.push_back(inputShape[i]);
          expandReassociation.push_back({expandedIdx});
          expandedIdx += 1;
        }
      }
      auto expandedType = RankedTensorType::get(expandedShape, elemType);
      tensorInput = tensor::ExpandShapeOp::create(ib, expandedType, tensorInput,
                                                  expandReassociation);

      // 3. Transpose to move nRanks (at position scatterDim) to position 0:
      //    {d0, ..., nRanks, o_sd, ..., dN} -> {nRanks, d0, ..., o_sd, ..., dN}
      SmallVector<int64_t> permutation, transposedShape;
      permutation.emplace_back(scatterDim);
      for (int64_t i = 0; i < scatterDim; ++i)
        permutation.emplace_back(i);
      for (int64_t i = scatterDim + 1; i < (int64_t)expandedShape.size(); ++i)
        permutation.emplace_back(i);
      for (auto p : permutation)
        transposedShape.emplace_back(expandedShape[p]);

      Value permOutput = tensor::EmptyOp::create(ib, transposedShape, elemType);
      tensorInput =
          linalg::TransposeOp::create(ib, tensorInput, permOutput, permutation)
              ->getResult(0);

      // 4. Materialize as contiguous memref for MPI by copying into a
      //    freshly allocated buffer.
      auto mpiInType = MemRefType::get(transposedShape, elemType);
      Value transposedBuf =
          bufferization::ToBufferOp::create(ib, mpiInType, tensorInput, true);
      mpiInput = memref::AllocOp::create(ib, mpiInType);
      linalg::CopyOp::create(ib, transposedBuf, mpiInput);
    }

    // Allocate output buffer.
    Value output = memref::AllocOp::create(ib, outType);
    // Create the MPI ReduceScatter operation.
    mpi::ReduceScatterBlockOp::create(
        ib, TypeRange(), mpiInput, output,
        getMPIReductionOp(adaptor.getReductionAttr()), comm);

    // If the destination is a tensor, cast it to a tensor.
    if (isa<RankedTensorType>(op.getType()))
      output =
          bufferization::ToTensorOp::create(ib, op.getType(), output, true);
    else if (scatterDim != 0) // Deallocate the temporary input buffer
      memref::DeallocOp::create(ib, mpiInput);
    // Notice: If this is called from tensor-world, then we assume an extra pass
    // will take care of deallocating the intermediate buffers.

    rewriter.replaceOp(op, output);
    return success();
  }
};

struct ConvertAllGatherOp : public CommOpPattern<AllGatherOp> {
  using CommOpPattern::CommOpPattern;

  // shard.allgather concatenates along a specified gather-axis.
  // mpi.allgather always concatenates along the first dimension and
  // there is no MPI operation that allows gathering along an arbitrary axis.
  // Hence, if gather-axis != 0, we need to permute the output buffer
  // accordingly.

  LogicalResult
  matchAndRewrite(AllGatherOp op, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    SymbolTableCollection symbolTableCollection;
    FailureOr<GridOp> gridOp = checkGrid(op, symbolTableCollection);
    if (failed(gridOp))
      return failure();

    ImplicitLocOpBuilder ib(op.getLoc(), rewriter);
    Value input = getAsMemref(adaptor.getInput(), ib, true);
    MemRefType inType = cast<MemRefType>(input.getType());
    MemRefType outType = getMemrefType(cast<ShapedType>(op.getType()));
    auto inputShape = inType.getShape();
    auto outputShape = outType.getShape();
    int64_t gatherAxis = adaptor.getGatherAxisAttr().getInt();
    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(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});
      int64_t numGatherDims = gatherShape.size();
      for (int64_t i = 2; i < numGatherDims; ++i)
        reassociation.push_back({i});
      finalOutput = memref::CollapseShapeOp::create(ib, outType, finalOutput,
                                                    reassociation);

      // 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;
      for (int i = 1; i <= gatherAxis; ++i) {
        outShapePermuted.emplace_back(gatherShape[i]);
        permutation.emplace_back(i);
      }
      outShapePermuted.emplace_back(gatherShape[0]);
      permutation.emplace_back(0);
      for (size_t i = gatherAxis + 1; i < gatherShape.size(); ++i) {
        outShapePermuted.emplace_back(gatherShape[i]);
        permutation.emplace_back(i);
      }
      Value permOutput = tensor::EmptyOp::create(ib, outShapePermuted,
                                                 outType.getElementType());
      finalOutput =
          linalg::TransposeOp::create(ib, finalOutput, permOutput, permutation)
              ->getResult(0);

      // 3. Collapse the output buffer from {dim0, ..., nRanks, gatherAxis, ...}
      // to {dim0, ..., nRanks*gatherAxis, ...}.
      SmallVector<ReassociationIndices> reassociation;
      for (int64_t i = 0; i < gatherAxis; ++i) {
        reassociation.push_back({i});
      }
      reassociation.push_back({gatherAxis, gatherAxis + 1});
      for (int64_t i = gatherAxis + 2; i < (int64_t)outShapePermuted.size();
           ++i) {
        reassociation.push_back({i});
      }
      auto outTType =
          RankedTensorType::get(outputShape, outType.getElementType());
      finalOutput = tensor::CollapseShapeOp::create(ib, outTType, finalOutput,
                                                    reassociation);

      // 4. Cast back to memref if needed.
      if (isa<MemRefType>(op.getType()))
        finalOutput =
            bufferization::ToBufferOp::create(ib, outType, finalOutput, false);
    }

    rewriter.replaceOp(op, finalOutput);
    return success();
  }
};

struct ConvertUpdateHaloOp : public OpConversionPattern<UpdateHaloOp> {
  using OpConversionPattern::OpConversionPattern;

  LogicalResult
  matchAndRewrite(UpdateHaloOp op, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {

    // The input/output memref is assumed to be in C memory order.
    // Halos are exchanged as 2 blocks per dimension (one for each side: down
    // and up). For each haloed dimension `d`, the exchanged blocks are
    // expressed as multi-dimensional subviews. The subviews include potential
    // halos of higher dimensions `dh > d`, no halos for the lower dimensions
    // `dl < d` and for dimension `d` the currently exchanged halo only.
    // By iterating form higher to lower dimensions this also updates the halos
    // in the 'corners'.
    // memref.subview is used to read and write the halo data from and to the
    // local data. Because subviews and halos can have mixed dynamic and static
    // shapes, OpFoldResults are used whenever possible.

    auto haloSizes = getMixedValues(adaptor.getStaticHaloSizes(),
                                    adaptor.getHaloSizes(), rewriter);
    if (haloSizes.empty()) {
      // no halos -> nothing to do
      rewriter.replaceOp(op, adaptor.getDestination());
      return success();
    }

    SymbolTableCollection symbolTableCollection;
    Location loc = op.getLoc();

    // convert a OpFoldResult into a Value
    auto toValue = [&rewriter, &loc](OpFoldResult &v) -> Value {
      if (auto value = dyn_cast<Value>(v))
        return value;
      return arith::ConstantOp::create(
          rewriter, loc,
          rewriter.getIndexAttr(
              cast<IntegerAttr>(cast<Attribute>(v)).getInt()));
    };

    auto dest = adaptor.getDestination();
    auto dstShape = cast<ShapedType>(dest.getType()).getShape();
    Value array = dest;
    if (isa<RankedTensorType>(array.getType())) {
      // If the destination is a memref, we need to cast it to a tensor
      auto mmemrefType = MemRefType::get(
          dstShape, cast<ShapedType>(array.getType()).getElementType());
      array =
          bufferization::ToBufferOp::create(rewriter, loc, mmemrefType, array);
    }
    auto rank = cast<ShapedType>(array.getType()).getRank();
    auto opSplitAxes = adaptor.getSplitAxes().getAxes();
    auto grid = adaptor.getGrid();
    auto gridOp = getGrid(op, symbolTableCollection);
    // subviews need Index values
    for (auto &sz : haloSizes) {
      if (auto value = dyn_cast<Value>(sz))
        sz = arith::IndexCastOp::create(rewriter, loc, rewriter.getIndexType(),
                                        value)
                 .getResult();
    }

    // most of the offset/size/stride data is the same for all dims
    SmallVector<OpFoldResult> offsets(rank, rewriter.getIndexAttr(0));
    SmallVector<OpFoldResult> strides(rank, rewriter.getIndexAttr(1));
    SmallVector<OpFoldResult> shape(rank), dimSizes(rank);
    auto currHaloDim = -1; // halo sizes are provided for split dimensions only
    // we need the actual shape to compute offsets and sizes
    for (auto i = 0; i < rank; ++i) {
      auto s = dstShape[i];
      if (ShapedType::isDynamic(s))
        shape[i] = memref::DimOp::create(rewriter, loc, array, s).getResult();
      else
        shape[i] = rewriter.getIndexAttr(s);

      if ((size_t)i < opSplitAxes.size() && !opSplitAxes[i].empty()) {
        ++currHaloDim;
        // the offsets for lower dim sstarts after their down halo
        offsets[i] = haloSizes[currHaloDim * 2];

        // prepare shape and offsets of highest dim's halo exchange
        Value _haloSz = arith::AddIOp::create(
            rewriter, loc, toValue(haloSizes[currHaloDim * 2]),
            toValue(haloSizes[currHaloDim * 2 + 1]));
        // the halo shape of lower dims exlude the halos
        dimSizes[i] =
            arith::SubIOp::create(rewriter, loc, toValue(shape[i]), _haloSz)
                .getResult();
      } else {
        dimSizes[i] = shape[i];
      }
    }

    auto tagAttr = rewriter.getI32IntegerAttr(91); // we just pick something
    auto tag = arith::ConstantOp::create(rewriter, loc, tagAttr);
    auto zeroAttr = rewriter.getI32IntegerAttr(0); // for detecting v<0
    auto zero = arith::ConstantOp::create(rewriter, loc, zeroAttr);

    SmallVector<Type> indexResultTypes(gridOp.getShape().size(),
                                       rewriter.getIndexType());
    auto myMultiIndex =
        ProcessMultiIndexOp::create(rewriter, loc, indexResultTypes, grid)
            .getResult();
    // traverse all split axes from high to low dim
    for (ssize_t dim = opSplitAxes.size() - 1; dim >= 0; --dim) {
      auto splitAxes = opSplitAxes[dim];
      if (splitAxes.empty())
        continue;
      assert(currHaloDim >= 0 && (size_t)currHaloDim < haloSizes.size() / 2);
      // Get the linearized ids of the neighbors (down and up) for the
      // given split
      auto tmp = NeighborsLinearIndicesOp::create(rewriter, loc, grid,
                                                  myMultiIndex, splitAxes)
                     .getResults();
      // MPI operates on i32...
      Value neighbourIDs[2] = {
          arith::IndexCastOp::create(rewriter, loc, rewriter.getI32Type(),
                                     tmp[0]),
          arith::IndexCastOp::create(rewriter, loc, rewriter.getI32Type(),
                                     tmp[1])};

      auto lowerRecvOffset = rewriter.getIndexAttr(0);
      auto lowerSendOffset = toValue(haloSizes[currHaloDim * 2]);
      auto upperRecvOffset =
          arith::SubIOp::create(rewriter, loc, toValue(shape[dim]),
                                toValue(haloSizes[currHaloDim * 2 + 1]));
      auto upperSendOffset = arith::SubIOp::create(
          rewriter, loc, upperRecvOffset, toValue(haloSizes[currHaloDim * 2]));

      Value commWorld = mpi::CommWorldOp::create(
          rewriter, loc, mpi::CommType::get(op->getContext()));

      // Make sure we send/recv in a way that does not lead to a dead-lock.
      // The current approach is by far not optimal, this should be at least
      // be a red-black pattern or using MPI_sendrecv.
      // Also, buffers should be re-used.
      // Still using temporary contiguous buffers for MPI communication...
      // Still yielding a "serialized" communication pattern...
      auto genSendRecv = [&](bool upperHalo) {
        auto orgOffset = offsets[dim];
        dimSizes[dim] = upperHalo ? haloSizes[currHaloDim * 2 + 1]
                                  : haloSizes[currHaloDim * 2];
        // Check if we need to send and/or receive
        // Processes on the grid borders have only one neighbor
        auto to = upperHalo ? neighbourIDs[0] : neighbourIDs[1];
        auto from = upperHalo ? neighbourIDs[1] : neighbourIDs[0];
        auto hasFrom = arith::CmpIOp::create(
            rewriter, loc, arith::CmpIPredicate::sge, from, zero);
        auto hasTo = arith::CmpIOp::create(rewriter, loc,
                                           arith::CmpIPredicate::sge, to, zero);
        auto buffer = memref::AllocOp::create(
            rewriter, loc, dimSizes,
            cast<ShapedType>(array.getType()).getElementType());
        // if has neighbor: copy halo data from array to buffer and send
        scf::IfOp::create(
            rewriter, loc, hasTo, [&](OpBuilder &builder, Location loc) {
              offsets[dim] = upperHalo ? OpFoldResult(lowerSendOffset)
                                       : OpFoldResult(upperSendOffset);
              auto subview = memref::SubViewOp::create(
                  builder, loc, array, offsets, dimSizes, strides);
              memref::CopyOp::create(builder, loc, subview, buffer);
              mpi::SendOp::create(builder, loc, TypeRange{}, buffer, tag, to,
                                  commWorld);
              scf::YieldOp::create(builder, loc);
            });
        // if has neighbor: receive halo data into buffer and copy to array
        scf::IfOp::create(
            rewriter, loc, hasFrom, [&](OpBuilder &builder, Location loc) {
              offsets[dim] = upperHalo ? OpFoldResult(upperRecvOffset)
                                       : OpFoldResult(lowerRecvOffset);
              mpi::RecvOp::create(builder, loc, TypeRange{}, buffer, tag, from,
                                  commWorld);
              auto subview = memref::SubViewOp::create(
                  builder, loc, array, offsets, dimSizes, strides);
              memref::CopyOp::create(builder, loc, buffer, subview);
              scf::YieldOp::create(builder, loc);
            });
        memref::DeallocOp::create(rewriter, loc, buffer);
        offsets[dim] = orgOffset;
      };

      auto doSendRecv = [&](int upOrDown) {
        OpFoldResult &v = haloSizes[currHaloDim * 2 + upOrDown];
        Value haloSz = dyn_cast<Value>(v);
        if (!haloSz)
          haloSz = arith::ConstantOp::create(
              rewriter, loc,
              rewriter.getI32IntegerAttr(
                  cast<IntegerAttr>(cast<Attribute>(v)).getInt()));
        auto hasSize = arith::CmpIOp::create(
            rewriter, loc, arith::CmpIPredicate::sgt, haloSz, zero);
        scf::IfOp::create(rewriter, loc, hasSize,
                          [&](OpBuilder &builder, Location loc) {
                            genSendRecv(upOrDown > 0);
                            scf::YieldOp::create(builder, loc);
                          });
      };

      doSendRecv(0);
      doSendRecv(1);

      // the shape for lower dims include higher dims' halos
      dimSizes[dim] = shape[dim];
      // -> the offset for higher dims is always 0
      offsets[dim] = rewriter.getIndexAttr(0);
      // on to next halo
      --currHaloDim;
    }

    if (isa<MemRefType>(op.getResult().getType())) {
      rewriter.replaceOp(op, array);
    } else {
      assert(isa<RankedTensorType>(op.getResult().getType()));
      rewriter.replaceOp(op, bufferization::ToTensorOp::create(
                                 rewriter, loc, op.getResult().getType(), array,
                                 /*restrict=*/true, /*writable=*/true));
    }
    return success();
  }
};

struct ConvertShardToMPIPass
    : public impl::ConvertShardToMPIPassBase<ConvertShardToMPIPass> {
  using Base::Base;

  /// Run the dialect converter on the module.
  void runOnOperation() override {
    auto *ctxt = &getContext();
    RewritePatternSet patterns(ctxt);
    ConversionTarget target(getContext());

    // Define a type converter to convert shard::ShardingType,
    // mostly for use in return operations.
    TypeConverter typeConverter;
    typeConverter.addConversion([](Type type) { return type; });

    // convert shard::ShardingType to a tuple of RankedTensorTypes
    typeConverter.addConversion(
        [](ShardingType type,
           SmallVectorImpl<Type> &results) -> std::optional<LogicalResult> {
          auto i16 = IntegerType::get(type.getContext(), 16);
          auto i64 = IntegerType::get(type.getContext(), 64);
          std::array<int64_t, 2> shp = {ShapedType::kDynamic,
                                        ShapedType::kDynamic};
          results.emplace_back(RankedTensorType::get(shp, i16));
          results.emplace_back(RankedTensorType::get(shp, i64)); // actually ?x2
          results.emplace_back(RankedTensorType::get(shp, i64));
          return success();
        });

    // To 'extract' components, a UnrealizedConversionCastOp is expected
    // to define the input
    typeConverter.addTargetMaterialization(
        [&](OpBuilder &builder, TypeRange resultTypes, ValueRange inputs,
            Location loc) {
          // Expecting a single input.
          if (inputs.size() != 1 || !isa<TupleType>(inputs[0].getType()))
            return SmallVector<Value>();
          auto castOp = inputs[0].getDefiningOp<UnrealizedConversionCastOp>();
          // Expecting an UnrealizedConversionCastOp.
          if (!castOp)
            return SmallVector<Value>();
          // Fill a vector with elements of the tuple/castOp.
          SmallVector<Value> results;
          for (auto oprnd : castOp.getInputs()) {
            if (!isa<RankedTensorType>(oprnd.getType()))
              return SmallVector<Value>();
            results.emplace_back(oprnd);
          }
          return results;
        });

    // No shard dialect should left after conversion...
    target.addIllegalDialect<shard::ShardDialect>();
    // ...except the global GridOp. GridShapeOp which will get folded later.
    target.addLegalOp<shard::GridOp, shard::GridShapeOp>();
    // Allow all the stuff that our patterns will convert to
    target.addLegalDialect<BuiltinDialect, mpi::MPIDialect, scf::SCFDialect,
                           arith::ArithDialect, tensor::TensorDialect,
                           bufferization::BufferizationDialect,
                           linalg::LinalgDialect, memref::MemRefDialect,
                           affine::AffineDialect, cf::ControlFlowDialect>();
    // Make sure the function signature, calls etc. are legal
    target.addDynamicallyLegalOp<func::FuncOp>([&](func::FuncOp op) {
      return typeConverter.isSignatureLegal(op.getFunctionType());
    });
    target.addDynamicallyLegalOp<func::CallOp, func::ReturnOp>(
        [&](Operation *op) { return typeConverter.isLegal(op); });

    patterns.add<ConvertUpdateHaloOp, ConvertNeighborsLinearIndicesOp,
                 ConvertGetShardingOp, ConvertShardingOp, ConvertShardShapeOp,
                 ConvertAllGatherOp, ConvertAllReduceOp, ConvertReduceScatterOp,
                 ConvertProcessLinearIndexOp>(typeConverter, ctxt);
    SymbolTableCollection stc;
    populateProcessMultiIndexOpLoweringPatterns(patterns, stc);
    populateAllSliceOpLoweringPatterns(patterns, stc);

    populateFunctionOpInterfaceTypeConversionPattern<func::FuncOp>(
        patterns, typeConverter);
    populateCallOpTypeConversionPattern(patterns, typeConverter);
    populateReturnOpTypeConversionPattern(patterns, typeConverter);

    (void)applyPartialConversion(getOperation(), target, std::move(patterns));

    // Folding patterns cannot be mixed with conversion patterns -> extra pass.
    patterns.clear();
    SymbolTableCollection symbolTableCollection;
    mlir::shard::populateFoldingPatterns(patterns, symbolTableCollection);
    (void)applyPatternsGreedily(getOperation(), std::move(patterns));
  }
};

} // namespace