MeshOps.cpp

Go to the documentation of this file.

//===- MeshOps.cpp - Mesh Dialect Operations ------------------------------===//
//
// 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
//
//===----------------------------------------------------------------------===//
 
#include "mlir/Dialect/Mesh/IR/MeshOps.h"
 
#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/Mesh/IR/MeshDialect.h"
#include "mlir/Dialect/Utils/StaticValueUtils.h"
#include "mlir/IR/Attributes.h"
#include "mlir/IR/BuiltinAttributes.h"
#include "mlir/IR/BuiltinTypeInterfaces.h"
#include "mlir/IR/BuiltinTypes.h"
#include "mlir/IR/Diagnostics.h"
#include "mlir/IR/DialectImplementation.h"
#include "mlir/IR/Location.h"
#include "mlir/IR/PatternMatch.h"
#include "mlir/IR/TypeUtilities.h"
#include "mlir/IR/Value.h"
#include "mlir/Interfaces/ViewLikeInterface.h"
#include "mlir/Support/LLVM.h"
#include "llvm/ADT/ArrayRef.h"
#include "llvm/ADT/STLExtras.h"
#include "llvm/ADT/SmallSet.h"
#include "llvm/ADT/SmallVector.h"
#include "llvm/ADT/TypeSwitch.h"
#include "llvm/Support/Casting.h"
#include <algorithm>
#include <functional>
#include <iterator>
#include <numeric>
#include <optional>
#include <utility>
 
#define DEBUG_TYPE "mesh-ops"
#define DBGS() (llvm::dbgs() << "[" DEBUG_TYPE << "]: ")
 
using namespace mlir;
using namespace mlir::mesh;
 
#include "mlir/Dialect/Mesh/IR/MeshDialect.cpp.inc"
 
namespace {
 
struct DimensionSize {
  static DimensionSize dynamic() { return DimensionSize(ShapedType::kDynamic); }
  DimensionSize(int64_t val) : val(val) {}
  int64_t value() const { return val; }
  operator int64_t() const { return val; }
  bool isDynamic() const { return ShapedType::isDynamic(val); }
 
private:
  int64_t val;
};
 
} // namespace
 
static DimensionSize operator/(DimensionSize lhs, DimensionSize rhs) {
  if (lhs.isDynamic() || rhs.isDynamic()) {
    return DimensionSize::dynamic();
  }
  return lhs.value() / rhs.value();
}
 
static DimensionSize operator*(DimensionSize lhs, DimensionSize rhs) {
  if (lhs.isDynamic() || rhs.isDynamic()) {
    return DimensionSize::dynamic();
  }
  return lhs.value() * rhs.value();
}
 
//===----------------------------------------------------------------------===//
// Mesh dialect
//===----------------------------------------------------------------------===//
 
void MeshDialect::initialize() {
  addOperations<
#define GET_OP_LIST
#include "mlir/Dialect/Mesh/IR/MeshOps.cpp.inc"
      >();
  addAttributes<
#define GET_ATTRDEF_LIST
#include "mlir/Dialect/Mesh/IR/MeshAttributes.cpp.inc"
      >();
}
 
Operation *MeshDialect::materializeConstant(OpBuilder &builder, Attribute value,
                                            Type type, Location loc) {
  return arith::ConstantOp::materialize(builder, value, type, loc);
}
 
//===----------------------------------------------------------------------===//
// Mesh utilities
//===----------------------------------------------------------------------===//
 
static FailureOr<MeshOp> getMeshAndVerify(Operation *op,
                                          FlatSymbolRefAttr meshSymbol,
                                          SymbolTableCollection &symbolTable) {
  mesh::MeshOp mesh = getMeshOrNull(op, meshSymbol, symbolTable);
  if (!mesh) {
    return op->emitError() << "Undefined required mesh symbol \""
                           << meshSymbol.getValue() << "\".";
  }
 
  return mesh;
}
 
template <typename It>
bool isUnique(It begin, It end) {
  if (begin == end) {
    return true;
  }
  It next = std::next(begin);
  if (next == end) {
    return true;
  }
  for (; next != end; ++next, ++begin) {
    if (*begin == *next) {
      return false;
    }
  }
  return true;
}
 
static LogicalResult verifyMeshAxes(Location loc, ArrayRef<MeshAxis> axes,
                                    MeshOp mesh) {
  SmallVector<MeshAxis> sorted = llvm::to_vector(axes);
  llvm::sort(sorted);
  if (!isUnique(sorted.begin(), sorted.end())) {
    return emitError(loc) << "Mesh axes contains duplicate elements.";
  }
 
  MeshAxis rank = mesh.getRank();
  for (auto axis : axes) {
    if (axis >= rank || axis < 0) {
      return emitError(loc)
             << "0-based mesh axis index " << axis
             << " is out of bounds. The referenced mesh \"" << mesh.getSymName()
             << "\" is of rank " << rank << ".";
    }
  }
 
  return success();
}
 
template <typename InShape, typename MeshShape, typename SplitAxes,
          typename OutShape>
static void shardShape(const InShape &inShape, const MeshShape &meshShape,
                       const SplitAxes &splitAxes, OutShape &outShape) {
  std::copy(llvm::adl_begin(inShape), llvm::adl_end(inShape),
            llvm::adl_begin(outShape));
  for (auto [tensorAxis, innerSplitAxes] : llvm::enumerate(splitAxes)) {
    outShape[tensorAxis] = shardDimension(
        inShape[tensorAxis],
        collectiveProcessGroupSize(innerSplitAxes.asArrayRef(), meshShape));
  }
}
 
ShapedType mesh::shardShapedType(ShapedType shape, MeshOp mesh,
                                 MeshShardingAttr sharding) {
  using Dim = std::decay_t<decltype(shape.getDimSize(0))>;
  SmallVector<Dim> resShapeArr(shape.getShape().size());
  shardShape(shape.getShape(), mesh.getShape(), sharding.getSplitAxes(),
             resShapeArr);
  return shape.clone(resShapeArr);
}
 
Type mesh::shardType(Type type, MeshOp mesh, MeshShardingAttr sharding) {
  RankedTensorType rankedTensorType = dyn_cast<RankedTensorType>(type);
  if (rankedTensorType) {
    return shardShapedType(rankedTensorType, mesh, sharding);
  }
 
  assert(!sharding);
  return type;
}
 
void mlir::mesh::maybeInsertTargetShardingAnnotation(MeshShardingAttr sharding,
                                                     OpOperand &operand,
                                                     OpBuilder &builder) {
  OpBuilder::InsertionGuard insertionGuard(builder);
  Value operandValue = operand.get();
  Operation *operandOp = operand.getOwner();
  builder.setInsertionPointAfterValue(operandValue);
  ShardOp shardOp = dyn_cast<ShardOp>(operandOp);
  if (shardOp && shardOp.getShard() == sharding &&
      !shardOp.getAnnotateForUsers()) {
    // No need for anything the correct sharding is already set.
    return;
  }
 
  auto newShardOp =
      builder.create<ShardOp>(operandValue.getLoc(), operandValue, sharding,
                              /*annotate_for_users*/ false);
  IRRewriter rewriter(builder);
  rewriter.replaceUsesWithIf(
      operandValue, newShardOp, [operandOp, operandValue](OpOperand &use) {
        return use.getOwner() == operandOp && use.get() == operandValue;
      });
 
  if (!shardOp || shardOp.getAnnotateForUsers()) {
    return;
  }
 
  auto newShardOp2 = builder.create<ShardOp>(
      operandValue.getLoc(), newShardOp, sharding, /*annotate_for_users*/ true);
  rewriter.replaceAllUsesExcept(newShardOp, newShardOp2, newShardOp2);
}
 
void mlir::mesh::maybeInsertTargetShardingAnnotation(MeshShardingAttr sharding,
                                                     OpResult result,
                                                     OpBuilder &builder) {
  for (auto &use : llvm::make_early_inc_range(result.getUses())) {
    maybeInsertTargetShardingAnnotation(sharding, use, builder);
  }
}
 
void mlir::mesh::maybeInsertSourceShardingAnnotation(MeshShardingAttr sharding,
                                                     OpOperand &operand,
                                                     OpBuilder &builder) {
  OpBuilder::InsertionGuard insertionGuard(builder);
  Value operandValue = operand.get();
  Operation *operandOp = operand.getOwner();
  Operation *operandSrcOp = operandValue.getDefiningOp();
  bool isBlockArg = !operandSrcOp;
  ShardOp shardOp = dyn_cast_or_null<ShardOp>(operandSrcOp);
 
  if (shardOp && shardOp.getShard() == sharding &&
      shardOp.getAnnotateForUsers()) {
    // No need for anything the correct sharding is already set.
    return;
  }
 
  builder.setInsertionPoint(operandOp);
  auto newShardOp =
      builder.create<ShardOp>(operandValue.getLoc(), operandValue, sharding,
                              /*annotate_for_users*/ true);
  IRRewriter rewriter(builder);
  rewriter.replaceUsesWithIf(
      operandValue, newShardOp, [operandOp, operandValue](OpOperand &use) {
        return use.getOwner() == operandOp && use.get() == operandValue;
      });
 
  if (isBlockArg || !shardOp || !shardOp.getAnnotateForUsers()) {
    // No need for resharding.
    return;
  }
 
  builder.setInsertionPoint(newShardOp);
  auto newPreceedingShardOp =
      builder.create<ShardOp>(operandValue.getLoc(), operandValue, sharding,
                              /*annotate_for_users*/ false);
  rewriter.replaceUsesWithIf(newShardOp.getOperand(), newPreceedingShardOp,
                             [&newShardOp](OpOperand &use) {
                               return use.getOwner() ==
                                      newShardOp.getOperation();
                             });
}
 
//===----------------------------------------------------------------------===//
// mesh.mesh op
//===----------------------------------------------------------------------===//
 
LogicalResult MeshOp::verify() {
  int64_t rank = getRank();
 
  if (rank <= 0)
    return emitOpError("rank of mesh is expected to be a positive integer");
 
  for (int64_t dimSize : getShape()) {
    if (dimSize < 0 && !ShapedType::isDynamic(dimSize))
      return emitOpError("dimension size of a mesh is expected to be "
                         "non-negative or dynamic");
  }
 
  return success();
}
 
//===----------------------------------------------------------------------===//
// mesh.mesh_shape op
//===----------------------------------------------------------------------===//
 
LogicalResult
MeshShapeOp::verifySymbolUses(SymbolTableCollection &symbolTable) {
  auto mesh = ::getMeshAndVerify(getOperation(), getMeshAttr(), symbolTable);
  if (failed(mesh)) {
    return failure();
  }
  if (failed(verifyMeshAxes(getLoc(), getAxes(), mesh.value()))) {
    return failure();
  }
 
  size_t expectedResultsCount =
      getAxes().empty() ? mesh->getRank() : getAxes().size();
  if (getResult().size() != expectedResultsCount) {
    return emitError() << "Unexpected number of results " << getResult().size()
                       << ". Expected " << expectedResultsCount << ".";
  }
 
  return success();
}
 
void MeshShapeOp::build(OpBuilder &odsBuilder, OperationState &odsState,
                        MeshOp mesh) {
  build(odsBuilder, odsState, mesh, SmallVector<MeshAxis>());
}
 
void MeshShapeOp::build(OpBuilder &odsBuilder, OperationState &odsState,
                        MeshOp mesh, ArrayRef<MeshAxis> axes) {
  build(odsBuilder, odsState,
        SmallVector<Type>(axes.empty() ? mesh.getRank() : axes.size(),
                          odsBuilder.getIndexType()),
        mesh.getSymName(), MeshAxesAttr::get(odsBuilder.getContext(), axes));
}
 
void MeshShapeOp::build(OpBuilder &odsBuilder, OperationState &odsState,
                        StringRef mesh, ArrayRef<MeshAxis> axes) {
  assert(!axes.empty());
  build(odsBuilder, odsState,
        SmallVector<Type>(axes.size(), odsBuilder.getIndexType()), mesh,
        MeshAxesAttr::get(odsBuilder.getContext(), axes));
}
 
void MeshShapeOp::getAsmResultNames(
    function_ref<void(Value, StringRef)> setNameFn) {
  setNameFn(getResults()[0], "mesh_shape");
}
 
//===----------------------------------------------------------------------===//
// mesh.shard attr
//===----------------------------------------------------------------------===//
 
LogicalResult
MeshShardingAttr::verify(function_ref<InFlightDiagnostic()> emitError,
                         FlatSymbolRefAttr, ArrayRef<MeshAxesAttr> splitAxes,
                         ArrayRef<MeshAxis> partialAxes, ReductionKind) {
  // TODO: At present mesh symbol ref is not verified. This is due to the
  // difficulty in fetching the corresponding symbol op based on an attribute.
 
  llvm::SmallSet<MeshAxis, 4> visitedAxes;
 
  auto checkMeshAxis = [&](ArrayRef<MeshAxis> axesArray) -> LogicalResult {
    for (MeshAxis axis : axesArray) {
      if (axis < 0)
        return emitError() << "mesh axis is expected to be non-negative";
      if (!visitedAxes.insert(axis).second)
        return emitError() << "mesh axis duplicated";
    }
    return success();
  };
 
  for (MeshAxesAttr subAxes : splitAxes) {
    ArrayRef<MeshAxis> subAxesArray = subAxes.asArrayRef();
    if (failed(checkMeshAxis(subAxesArray)))
      return failure();
  }
  if (failed(checkMeshAxis(partialAxes)))
    return failure();
  return success();
}
 
bool MeshShardingAttr::operator==(Attribute rhs) const {
  MeshShardingAttr rhsAsMeshShardingAttr =
      mlir::dyn_cast<MeshShardingAttr>(rhs);
  return rhsAsMeshShardingAttr && *this == rhsAsMeshShardingAttr;
}
 
bool MeshShardingAttr::operator!=(Attribute rhs) const {
  return !(*this == rhs);
}
 
bool MeshShardingAttr::operator==(MeshShardingAttr rhs) const {
  if (getMesh() != rhs.getMesh() || getPartialAxes() != rhs.getPartialAxes()) {
    return false;
  }
 
  if (!getPartialAxes().empty() && getPartialType() != rhs.getPartialType()) {
    return false;
  }
 
  auto minSize = std::min(getSplitAxes().size(), rhs.getSplitAxes().size());
  if (!llvm::equal(llvm::make_range(getSplitAxes().begin(),
                                    getSplitAxes().begin() + minSize),
                   llvm::make_range(rhs.getSplitAxes().begin(),
                                    rhs.getSplitAxes().begin() + minSize))) {
    return false;
  }
 
  return llvm::all_of(llvm::make_range(getSplitAxes().begin() + minSize,
                                       getSplitAxes().end()),
                      std::mem_fn(&MeshAxesAttr::empty)) &&
         llvm::all_of(llvm::make_range(rhs.getSplitAxes().begin() + minSize,
                                       rhs.getSplitAxes().end()),
                      std::mem_fn(&MeshAxesAttr::empty));
}
 
bool MeshShardingAttr::operator!=(MeshShardingAttr rhs) const {
  return !(*this == rhs);
}
 
//===----------------------------------------------------------------------===//
// mesh.shard op
//===----------------------------------------------------------------------===//
 
void ShardOp::getAsmResultNames(
    function_ref<void(Value, StringRef)> setNameFn) {
  setNameFn(getResult(), "sharding_annotated");
}
 
//===----------------------------------------------------------------------===//
// mesh.process_multi_index op
//===----------------------------------------------------------------------===//
 
LogicalResult
ProcessMultiIndexOp::verifySymbolUses(SymbolTableCollection &symbolTable) {
  auto mesh = ::getMeshAndVerify(getOperation(), getMeshAttr(), symbolTable);
  if (failed(mesh)) {
    return failure();
  }
  if (failed(verifyMeshAxes(getLoc(), getAxes(), mesh.value()))) {
    return failure();
  }
 
  size_t expectedResultsCount =
      getAxes().empty() ? mesh->getRank() : getAxes().size();
  if (getResult().size() != expectedResultsCount) {
    return emitError() << "Unexpected number of results " << getResult().size()
                       << ". Expected " << expectedResultsCount << ".";
  }
 
  return success();
}
 
void ProcessMultiIndexOp::build(OpBuilder &odsBuilder, OperationState &odsState,
                                MeshOp mesh) {
  build(odsBuilder, odsState,
        SmallVector<Type>(mesh.getRank(), odsBuilder.getIndexType()),
        mesh.getSymName(), ArrayRef<MeshAxis>());
}
 
void ProcessMultiIndexOp::build(OpBuilder &odsBuilder, OperationState &odsState,
                                StringRef mesh, ArrayRef<MeshAxis> axes) {
  build(odsBuilder, odsState,
        SmallVector<Type>(axes.size(), odsBuilder.getIndexType()), mesh,
        MeshAxesAttr::get(odsBuilder.getContext(), axes));
}
 
void ProcessMultiIndexOp::getAsmResultNames(
    function_ref<void(Value, StringRef)> setNameFn) {
  setNameFn(getResults()[0], "proc_linear_idx");
}
 
//===----------------------------------------------------------------------===//
// mesh.process_linear_index op
//===----------------------------------------------------------------------===//
 
LogicalResult
ProcessLinearIndexOp::verifySymbolUses(SymbolTableCollection &symbolTable) {
  auto mesh = ::getMeshAndVerify(getOperation(), getMeshAttr(), symbolTable);
  if (failed(mesh)) {
    return failure();
  }
  return success();
}
 
void ProcessLinearIndexOp::build(OpBuilder &odsBuilder,
                                 OperationState &odsState, MeshOp mesh) {
  build(odsBuilder, odsState, mesh.getSymName());
}
 
void ProcessLinearIndexOp::getAsmResultNames(
    function_ref<void(Value, StringRef)> setNameFn) {
  setNameFn(getResult(), "proc_linear_idx");
}
 
//===----------------------------------------------------------------------===//
// collective communication ops
//===----------------------------------------------------------------------===//
 
namespace {
 
template <typename Op>
struct EmptyMeshAxesCanonicalizationPattern : OpRewritePattern<Op> {
  using OpRewritePattern<Op>::OpRewritePattern;
  LogicalResult matchAndRewrite(Op op,
                                PatternRewriter &rewriter) const override {
    auto meshAxes = op.getMeshAxes();
    if (!meshAxes.empty()) {
      return failure();
    }
    if (op.getInput().getType() != op.getResult().getType()) {
      return failure();
    }
 
    rewriter.replaceAllUsesWith(op.getResult(), op.getInput());
    rewriter.eraseOp(op.getOperation());
    return success();
  }
};
 
} // namespace
 
static LogicalResult verifyInGroupDevice(Location loc, StringRef deviceName,
                                         ArrayRef<int64_t> device,
                                         Operation::operand_range deviceDynamic,
                                         ArrayRef<MeshAxis> meshAxes,
                                         ArrayRef<int64_t> meshShape) {
  if (device.size() != meshAxes.size()) {
    return emitError(loc) << "In-group device \"" << deviceName
                          << "\" has unexpected multi-index size "
                          << device.size() << ". Expected " << meshAxes.size()
                          << ".";
  }
 
  for (size_t i = 0; i < device.size(); ++i) {
    if (!ShapedType::isDynamic(device[i]) &&
        !ShapedType::isDynamic(meshShape[meshAxes[i]]) &&
        meshShape[meshAxes[i]] <= device[i]) {
      return emitError(loc)
             << "Out of bounds coordinate " << i << " for in-group device \""
             << deviceName << "\"."
             << " Got " << device[i] << ", but expected value in the range [0, "
             << (meshShape[meshAxes[i]] - 1) << "].";
    }
  }
  return success();
}
 
template <typename Op>
static FailureOr<MeshOp>
getMeshAndVerifyAxes(Op op, SymbolTableCollection &symbolTable) {
  auto mesh =
      ::getMeshAndVerify(op.getOperation(), op.getMeshAttr(), symbolTable);
  if (failed(mesh)) {
    return failure();
  }
  if (failed(verifyMeshAxes(op.getLoc(), op.getMeshAxes(), mesh.value()))) {
    return failure();
  }
  return mesh;
}
 
template <typename It>
static auto product(It begin, It end) {
  using ElementType = std::decay_t<decltype(*begin)>;
  return std::accumulate(begin, end, static_cast<ElementType>(1),
                         std::multiplies<ElementType>());
}
 
template <typename R>
static auto product(R &&range) {
  return product(adl_begin(range), adl_end(range));
}
 
static LogicalResult verifyDimensionCompatibility(Location loc,
                                                  int64_t expectedDimSize,
                                                  int64_t resultDimSize,
                                                  int64_t resultAxis) {
  if (!ShapedType::isDynamic(resultDimSize) &&
      expectedDimSize != resultDimSize) {
    return emitError(loc) << "Dimension size mismatch for result axis "
                          << resultAxis << ". Expected "
                          << (ShapedType::isDynamic(expectedDimSize)
                                  ? Twine("dynamic")
                                  : Twine(expectedDimSize))
                          << ", but got " << resultDimSize << ".";
  }
 
  return success();
}
 
static LogicalResult verifyGatherOperandAndResultShape(
    Value operand, Value result, int64_t gatherAxis,
    ArrayRef<MeshAxis> meshAxes, ArrayRef<int64_t> meshShape) {
  auto resultRank = cast<ShapedType>(result.getType()).getRank();
  if (gatherAxis < 0 || gatherAxis >= resultRank) {
    return emitError(result.getLoc())
           << "Gather axis " << gatherAxis << " is out of bounds [0, "
           << resultRank << ").";
  }
 
  ShapedType operandType = cast<ShapedType>(operand.getType());
  ShapedType resultType = cast<ShapedType>(result.getType());
  auto deviceGroupSize =
      DimensionSize(collectiveProcessGroupSize(meshAxes, meshShape));
  for (int64_t axis = 0; axis < operandType.getRank(); ++axis) {
    auto operandDimSize = DimensionSize(operandType.getDimSize(axis));
    auto resultDimSize = DimensionSize(resultType.getDimSize(axis));
    auto expectedResultDimSize =
        axis == gatherAxis ? deviceGroupSize * operandDimSize : operandDimSize;
    if (failed(verifyDimensionCompatibility(
            result.getLoc(), expectedResultDimSize, resultDimSize, axis))) {
      return failure();
    }
  }
  return success();
}
 
static LogicalResult verifyAllToAllOperandAndResultShape(
    Value operand, Value result, int64_t splitAxis, int64_t concatAxis,
    ArrayRef<MeshAxis> meshAxes, ArrayRef<int64_t> meshShape) {
  ShapedType operandType = cast<ShapedType>(operand.getType());
  ShapedType resultType = cast<ShapedType>(result.getType());
  for (int64_t axis = 0; axis < operandType.getRank(); ++axis) {
    if ((axis != splitAxis && axis != concatAxis) || splitAxis == concatAxis) {
      if (failed(verifyDimensionCompatibility(
              result.getLoc(), operandType.getDimSize(axis),
              resultType.getDimSize(axis), axis))) {
        return failure();
      }
    }
  }
 
  if (splitAxis == concatAxis) {
    return success();
  }
 
  auto deviceGroupSize =
      DimensionSize(collectiveProcessGroupSize(meshAxes, meshShape));
  auto operandConcatDimSize = DimensionSize(operandType.getDimSize(concatAxis));
  auto operandSplitDimSize = DimensionSize(operandType.getDimSize(splitAxis));
  DimensionSize expectedResultConcatDimSize =
      operandConcatDimSize * deviceGroupSize;
  DimensionSize expectedResultSplitDimSize =
      operandSplitDimSize / deviceGroupSize;
  if (!expectedResultSplitDimSize.isDynamic() &&
      int64_t(operandSplitDimSize) % int64_t(deviceGroupSize) != 0) {
    expectedResultSplitDimSize = DimensionSize::dynamic();
  }
  if (failed(verifyDimensionCompatibility(
          result.getLoc(), expectedResultConcatDimSize.value(),
          resultType.getDimSize(concatAxis), concatAxis))) {
    return failure();
  }
  if (failed(verifyDimensionCompatibility(
          result.getLoc(), expectedResultSplitDimSize.value(),
          resultType.getDimSize(splitAxis), splitAxis))) {
    return failure();
  }
 
  return success();
}
 
static LogicalResult verifyScatterOrSliceOperandAndResultShape(
    Value operand, Value result, int64_t tensorAxis,
    ArrayRef<MeshAxis> meshAxes, ArrayRef<int64_t> meshShape) {
  ShapedType operandType = cast<ShapedType>(operand.getType());
  ShapedType resultType = cast<ShapedType>(result.getType());
  for (int64_t axis = 0; axis < operandType.getRank(); ++axis) {
    if (axis != tensorAxis) {
      if (failed(verifyDimensionCompatibility(
              result.getLoc(), operandType.getDimSize(axis),
              resultType.getDimSize(axis), axis))) {
        return failure();
      }
    }
  }
 
  auto deviceGroupSize =
      DimensionSize(collectiveProcessGroupSize(meshAxes, meshShape));
  auto operandScatterDimSize =
      DimensionSize(operandType.getDimSize(tensorAxis));
  if (!operandScatterDimSize.isDynamic() && !deviceGroupSize.isDynamic() &&
      int64_t(operandScatterDimSize) % int64_t(deviceGroupSize) != 0) {
    return emitError(result.getLoc())
           << "Operand dimension size " << int64_t(operandScatterDimSize)
           << " is not divisible by collective device group size "
           << int64_t(deviceGroupSize) << " for tensor axis " << tensorAxis
           << ".";
  }
  DimensionSize expectedResultTensorDimSize =
      operandScatterDimSize / deviceGroupSize;
  if (failed(verifyDimensionCompatibility(
          result.getLoc(), expectedResultTensorDimSize.value(),
          resultType.getDimSize(tensorAxis), tensorAxis))) {
    return failure();
  }
 
  return success();
}
 
static RankedTensorType sliceResultType(Type operandType, MeshOp mesh,
                                        ArrayRef<MeshAxis> meshAxes,
                                        int64_t sliceAxis) {
  RankedTensorType operandRankedTensorType =
      cast<RankedTensorType>(operandType);
  DimensionSize operandSliceAxisSize =
      operandRankedTensorType.getShape()[sliceAxis];
  SmallVector<int64_t> resultShape =
      llvm::to_vector(operandRankedTensorType.getShape());
 
  resultShape[sliceAxis] =
      operandSliceAxisSize /
      DimensionSize(collectiveProcessGroupSize(meshAxes, mesh));
  return operandRankedTensorType.clone(resultShape);
}
 
//===----------------------------------------------------------------------===//
// mesh.all_gather op
//===----------------------------------------------------------------------===//
 
LogicalResult
AllGatherOp::verifySymbolUses(SymbolTableCollection &symbolTable) {
  auto mesh = getMeshAndVerifyAxes(*this, symbolTable);
  if (failed(mesh)) {
    return failure();
  }
  auto gatherAxis = getGatherAxis().getSExtValue();
  return verifyGatherOperandAndResultShape(getOperand(), getResult(),
                                           gatherAxis, getMeshAxes(),
                                           mesh.value().getShape());
}
 
void AllGatherOp::getCanonicalizationPatterns(RewritePatternSet &patterns,
                                              MLIRContext *context) {
  patterns.add<EmptyMeshAxesCanonicalizationPattern<AllGatherOp>>(context);
}
 
void AllGatherOp::getAsmResultNames(
    function_ref<void(Value, StringRef)> setNameFn) {
  setNameFn(getResult(), "all_gather");
}
 
//===----------------------------------------------------------------------===//
// mesh.all_reduce op
//===----------------------------------------------------------------------===//
 
LogicalResult
AllReduceOp::verifySymbolUses(SymbolTableCollection &symbolTable) {
  return getMeshAndVerifyAxes(*this, symbolTable);
}
 
void AllReduceOp::getCanonicalizationPatterns(RewritePatternSet &patterns,
                                              MLIRContext *context) {
  patterns.add<EmptyMeshAxesCanonicalizationPattern<AllReduceOp>>(context);
}
 
void AllReduceOp::build(OpBuilder &odsBuilder, OperationState &odsState,
                        Value input, StringRef mesh,
                        ArrayRef<MeshAxis> meshAxes, ReductionKind reduction) {
  build(odsBuilder, odsState, input.getType(), mesh, meshAxes, input,
        reduction);
}
 
void AllReduceOp::getAsmResultNames(
    function_ref<void(Value, StringRef)> setNameFn) {
  setNameFn(getResult(), "all_reduce");
}
 
//===----------------------------------------------------------------------===//
// mesh.all_slice op
//===----------------------------------------------------------------------===//
 
LogicalResult AllSliceOp::verifySymbolUses(SymbolTableCollection &symbolTable) {
  auto mesh = getMeshAndVerifyAxes(*this, symbolTable);
  if (failed(mesh)) {
    return failure();
  }
  return verifyScatterOrSliceOperandAndResultShape(
      getOperand(), getResult(), getSliceAxis().getSExtValue(), getMeshAxes(),
      mesh.value().getShape());
}
 
void AllSliceOp::getCanonicalizationPatterns(RewritePatternSet &patterns,
                                             MLIRContext *context) {
  patterns.add<EmptyMeshAxesCanonicalizationPattern<AllSliceOp>>(context);
}
 
void AllSliceOp::build(OpBuilder &odsBuilder, OperationState &odsState,
                       Value input, MeshOp mesh, ArrayRef<MeshAxis> meshAxes,
                       int64_t sliceAxis) {
  Type resultType = sliceResultType(input.getType(), mesh, meshAxes, sliceAxis);
  build(odsBuilder, odsState, resultType, input, mesh.getSymName(), meshAxes,
        sliceAxis);
}
 
void AllSliceOp::build(OpBuilder &odsBuilder, OperationState &odsState,
                       Type resultType, Value input, StringRef mesh,
                       ArrayRef<MeshAxis> meshAxes, int64_t sliceAxis) {
  build(odsBuilder, odsState, resultType, mesh, meshAxes, input,
        APInt(sizeof(sliceAxis) * CHAR_BIT, sliceAxis));
}
 
void AllSliceOp::getAsmResultNames(
    function_ref<void(Value, StringRef)> setNameFn) {
  setNameFn(getResult(), "all_slice");
}
 
//===----------------------------------------------------------------------===//
// mesh.all_to_all op
//===----------------------------------------------------------------------===//
 
LogicalResult AllToAllOp::verifySymbolUses(SymbolTableCollection &symbolTable) {
  auto mesh = getMeshAndVerifyAxes(*this, symbolTable);
  if (failed(mesh)) {
    return failure();
  }
 
  return verifyAllToAllOperandAndResultShape(
      getOperand(), getResult(), getSplitAxis().getSExtValue(),
      getConcatAxis().getSExtValue(), getMeshAxes(), mesh.value().getShape());
}
 
void AllToAllOp::getCanonicalizationPatterns(RewritePatternSet &patterns,
                                             MLIRContext *context) {
  patterns.add<EmptyMeshAxesCanonicalizationPattern<AllToAllOp>>(context);
}
 
void AllToAllOp::getAsmResultNames(
    function_ref<void(Value, StringRef)> setNameFn) {
  setNameFn(getResult(), "all_to_all");
}
 
//===----------------------------------------------------------------------===//
// mesh.broadcast op
//===----------------------------------------------------------------------===//
 
LogicalResult
BroadcastOp::verifySymbolUses(SymbolTableCollection &symbolTable) {
  auto mesh = getMeshAndVerifyAxes(*this, symbolTable);
  if (failed(mesh)) {
    return failure();
  }
  if (failed(verifyInGroupDevice(getLoc(), getRootAttrName(), getRoot(),
                                 getRootDynamic(), getMeshAxes(),
                                 mesh.value().getShape()))) {
    return failure();
  }
 
  return success();
}
 
void BroadcastOp::getCanonicalizationPatterns(RewritePatternSet &patterns,
                                              MLIRContext *context) {
  patterns.add<EmptyMeshAxesCanonicalizationPattern<BroadcastOp>>(context);
}
 
void BroadcastOp::getAsmResultNames(
    function_ref<void(Value, StringRef)> setNameFn) {
  setNameFn(getResult(), "broadcast");
}
 
//===----------------------------------------------------------------------===//
// mesh.gather op
//===----------------------------------------------------------------------===//
 
LogicalResult GatherOp::verifySymbolUses(SymbolTableCollection &symbolTable) {
  auto mesh = getMeshAndVerifyAxes(*this, symbolTable);
  if (failed(mesh)) {
    return failure();
  }
  if (failed(verifyInGroupDevice(getLoc(), getRootAttrName(), getRoot(),
                                 getRootDynamic(), getMeshAxes(),
                                 mesh.value().getShape()))) {
    return failure();
  }
 
  auto gatherAxis = getGatherAxis().getSExtValue();
  return verifyGatherOperandAndResultShape(getInput(), getResult(), gatherAxis,
                                           getMeshAxes(),
                                           mesh.value().getShape());
}
 
void GatherOp::getCanonicalizationPatterns(RewritePatternSet &patterns,
                                           MLIRContext *context) {
  patterns.add<EmptyMeshAxesCanonicalizationPattern<GatherOp>>(context);
}
 
void GatherOp::getAsmResultNames(
    function_ref<void(Value, StringRef)> setNameFn) {
  setNameFn(getResult(), "gather");
}
 
//===----------------------------------------------------------------------===//
// mesh.recv op
//===----------------------------------------------------------------------===//
 
LogicalResult RecvOp::verifySymbolUses(SymbolTableCollection &symbolTable) {
  auto mesh = getMeshAndVerifyAxes(*this, symbolTable);
  if (failed(mesh)) {
    return failure();
  }
  if (getSource() &&
      failed(verifyInGroupDevice(getLoc(), getSourceAttrName(),
                                 getSource().value(), getSourceDynamic(),
                                 getMeshAxes(), mesh.value().getShape()))) {
    return failure();
  }
  return success();
}
 
void RecvOp::getCanonicalizationPatterns(RewritePatternSet &patterns,
                                         MLIRContext *context) {
  patterns.add<EmptyMeshAxesCanonicalizationPattern<RecvOp>>(context);
}
 
void RecvOp::getAsmResultNames(function_ref<void(Value, StringRef)> setNameFn) {
  setNameFn(getResult(), "recv");
}
 
//===----------------------------------------------------------------------===//
// mesh.reduce op
//===----------------------------------------------------------------------===//
 
LogicalResult ReduceOp::verifySymbolUses(SymbolTableCollection &symbolTable) {
  auto mesh = getMeshAndVerifyAxes(*this, symbolTable);
  if (failed(mesh)) {
    return failure();
  }
  if (failed(verifyInGroupDevice(getLoc(), getRootAttrName(), getRoot(),
                                 getRootDynamic(), getMeshAxes(),
                                 mesh.value().getShape()))) {
    return failure();
  }
 
  return success();
}
 
void ReduceOp::getCanonicalizationPatterns(RewritePatternSet &patterns,
                                           MLIRContext *context) {
  patterns.add<EmptyMeshAxesCanonicalizationPattern<ReduceOp>>(context);
}
 
void ReduceOp::getAsmResultNames(
    function_ref<void(Value, StringRef)> setNameFn) {
  setNameFn(getResult(), "reduce");
}
 
//===----------------------------------------------------------------------===//
// mesh.reduce_scatter op
//===----------------------------------------------------------------------===//
 
LogicalResult
ReduceScatterOp::verifySymbolUses(SymbolTableCollection &symbolTable) {
  auto mesh = getMeshAndVerifyAxes(*this, symbolTable);
  if (failed(mesh)) {
    return failure();
  }
 
  return verifyScatterOrSliceOperandAndResultShape(
      getOperand(), getResult(), getScatterAxis().getSExtValue(), getMeshAxes(),
      mesh.value().getShape());
}
 
void ReduceScatterOp::getCanonicalizationPatterns(RewritePatternSet &patterns,
                                                  MLIRContext *context) {
  patterns.add<EmptyMeshAxesCanonicalizationPattern<ReduceScatterOp>>(context);
}
 
void ReduceScatterOp::getAsmResultNames(
    function_ref<void(Value, StringRef)> setNameFn) {
  setNameFn(getResult(), "reduce_scatter");
}
 
//===----------------------------------------------------------------------===//
// mesh.scatter op
//===----------------------------------------------------------------------===//
 
LogicalResult ScatterOp::verifySymbolUses(SymbolTableCollection &symbolTable) {
  auto mesh = getMeshAndVerifyAxes(*this, symbolTable);
  if (failed(mesh)) {
    return failure();
  }
  if (failed(verifyInGroupDevice(getLoc(), getRootAttrName(), getRoot(),
                                 getRootDynamic(), getMeshAxes(),
                                 mesh.value().getShape()))) {
    return failure();
  }
 
  auto scatterAxis = getScatterAxis().getSExtValue();
  return verifyScatterOrSliceOperandAndResultShape(getInput(), getResult(),
                                                   scatterAxis, getMeshAxes(),
                                                   mesh.value().getShape());
}
 
void ScatterOp::getCanonicalizationPatterns(RewritePatternSet &patterns,
                                            MLIRContext *context) {
  patterns.add<EmptyMeshAxesCanonicalizationPattern<ScatterOp>>(context);
}
 
void ScatterOp::getAsmResultNames(
    function_ref<void(Value, StringRef)> setNameFn) {
  setNameFn(getResult(), "scatter");
}
 
//===----------------------------------------------------------------------===//
// mesh.send op
//===----------------------------------------------------------------------===//
 
LogicalResult SendOp::verifySymbolUses(SymbolTableCollection &symbolTable) {
  auto mesh = getMeshAndVerifyAxes(*this, symbolTable);
  if (failed(mesh)) {
    return failure();
  }
  if (failed(verifyInGroupDevice(getLoc(), getDestinationAttrName(),
                                 getDestination(), getDestinationDynamic(),
                                 getMeshAxes(), mesh.value().getShape()))) {
    return failure();
  }
  return success();
}
 
void SendOp::getCanonicalizationPatterns(RewritePatternSet &patterns,
                                         MLIRContext *context) {
  patterns.add<EmptyMeshAxesCanonicalizationPattern<SendOp>>(context);
}
 
void SendOp::getAsmResultNames(function_ref<void(Value, StringRef)> setNameFn) {
  setNameFn(getResult(), "send");
}
 
//===----------------------------------------------------------------------===//
// mesh.shift op
//===----------------------------------------------------------------------===//
 
LogicalResult ShiftOp::verifySymbolUses(SymbolTableCollection &symbolTable) {
  auto mesh = getMeshAndVerifyAxes(*this, symbolTable);
  if (failed(mesh)) {
    return failure();
  }
 
  auto meshAxes = getMeshAxes();
  auto shiftAxis = getShiftAxis().getZExtValue();
  if (llvm::find(meshAxes, shiftAxis) == meshAxes.end()) {
    return emitError() << "Invalid shift axis " << shiftAxis
                       << ". It must be one of the grouping mesh axes.";
  }
 
  return success();
}
 
void ShiftOp::getCanonicalizationPatterns(RewritePatternSet &patterns,
                                          MLIRContext *context) {
  // TODO: remove op when offset is 0 or if it is a rotate with and
  // offset % shift_axis_mesh_dim_size == 0.
}
 
void ShiftOp::getAsmResultNames(
    function_ref<void(Value, StringRef)> setNameFn) {
  setNameFn(getResult(), "shift");
}
 
//===----------------------------------------------------------------------===//
// TableGen'd op method definitions
//===----------------------------------------------------------------------===//
 
#define GET_OP_CLASSES
#include "mlir/Dialect/Mesh/IR/MeshOps.cpp.inc"
 
#define GET_ATTRDEF_CLASSES
#include "mlir/Dialect/Mesh/IR/MeshAttributes.cpp.inc"
 
#include "mlir/Dialect/Mesh/IR/MeshEnums.cpp.inc"