MLIR  21.0.0git
MeshOps.cpp
Go to the documentation of this file.
1 //===- MeshOps.cpp - Mesh Dialect Operations ------------------------------===//
2 //
3 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4 // See https://llvm.org/LICENSE.txt for license information.
5 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6 //
7 //===----------------------------------------------------------------------===//
8 
10 
14 #include "mlir/IR/Attributes.h"
17 #include "mlir/IR/BuiltinTypes.h"
18 #include "mlir/IR/Diagnostics.h"
20 #include "mlir/IR/IRMapping.h"
21 #include "mlir/IR/Location.h"
22 #include "mlir/IR/PatternMatch.h"
23 #include "mlir/IR/TypeUtilities.h"
24 #include "mlir/IR/Value.h"
26 #include "mlir/Support/LLVM.h"
28 #include "llvm/ADT/ArrayRef.h"
29 #include "llvm/ADT/STLExtras.h"
30 #include "llvm/ADT/SmallSet.h"
31 #include "llvm/ADT/SmallVector.h"
32 #include "llvm/ADT/TypeSwitch.h"
33 #include "llvm/Support/Casting.h"
34 #include <algorithm>
35 #include <functional>
36 #include <iterator>
37 #include <numeric>
38 #include <optional>
39 #include <utility>
40 
41 #define DEBUG_TYPE "mesh-ops"
42 #define DBGS() (llvm::dbgs() << "[" DEBUG_TYPE << "]: ")
43 
44 using namespace mlir;
45 using namespace mlir::mesh;
46 
47 #include "mlir/Dialect/Mesh/IR/MeshDialect.cpp.inc"
48 
49 namespace {
50 
51 struct DimensionSize {
52  static DimensionSize dynamic() { return DimensionSize(ShapedType::kDynamic); }
53  DimensionSize(int64_t val) : val(val) {}
54  int64_t value() const { return val; }
55  operator int64_t() const { return val; }
56  bool isDynamic() const { return ShapedType::isDynamic(val); }
57 
58 private:
59  int64_t val;
60 };
61 
62 } // namespace
63 
64 static DimensionSize operator/(DimensionSize lhs, DimensionSize rhs) {
65  if (lhs.isDynamic() || rhs.isDynamic()) {
66  return DimensionSize::dynamic();
67  }
68  return lhs.value() / rhs.value();
69 }
70 
71 static DimensionSize operator*(DimensionSize lhs, DimensionSize rhs) {
72  if (lhs.isDynamic() || rhs.isDynamic()) {
73  return DimensionSize::dynamic();
74  }
75  return lhs.value() * rhs.value();
76 }
77 
78 //===----------------------------------------------------------------------===//
79 // Inliner
80 //===----------------------------------------------------------------------===//
81 
82 namespace {
83 struct MeshInlinerInterface : public DialectInlinerInterface {
85  // Currently no restrictions are encoded for inlining.
86  bool isLegalToInline(Operation *, Operation *, bool) const final {
87  return true;
88  }
89  bool isLegalToInline(Region *, Region *, bool, IRMapping &) const final {
90  return true;
91  }
92  bool isLegalToInline(Operation *, Region *, bool, IRMapping &) const final {
93  return true;
94  }
95 };
96 } // namespace
97 
98 //===----------------------------------------------------------------------===//
99 // Mesh dialect
100 //===----------------------------------------------------------------------===//
101 
102 void MeshDialect::initialize() {
103  addOperations<
104 #define GET_OP_LIST
105 #include "mlir/Dialect/Mesh/IR/MeshOps.cpp.inc"
106  >();
107  addAttributes<
108 #define GET_ATTRDEF_LIST
109 #include "mlir/Dialect/Mesh/IR/MeshAttributes.cpp.inc"
110  >();
111  addTypes<
112 #define GET_TYPEDEF_LIST
113 #include "mlir/Dialect/Mesh/IR/MeshTypes.cpp.inc"
114  >();
115  addInterface<MeshInlinerInterface>();
116 }
117 
119  Type type, Location loc) {
120  return arith::ConstantOp::materialize(builder, value, type, loc);
121 }
122 
123 //===----------------------------------------------------------------------===//
124 // Mesh utilities
125 //===----------------------------------------------------------------------===//
126 
127 static FailureOr<MeshOp> getMeshAndVerify(Operation *op,
128  FlatSymbolRefAttr meshSymbol,
129  SymbolTableCollection &symbolTable) {
130  mesh::MeshOp mesh = getMeshOrNull(op, meshSymbol, symbolTable);
131  if (!mesh) {
132  return op->emitError() << "Undefined required mesh symbol \""
133  << meshSymbol.getValue() << "\".";
134  }
135 
136  return mesh;
137 }
138 
139 template <typename It>
140 bool isUnique(It begin, It end) {
141  if (begin == end) {
142  return true;
143  }
144  It next = std::next(begin);
145  if (next == end) {
146  return true;
147  }
148  for (; next != end; ++next, ++begin) {
149  if (*begin == *next) {
150  return false;
151  }
152  }
153  return true;
154 }
155 
156 static LogicalResult verifyMeshAxes(Location loc, ArrayRef<MeshAxis> axes,
157  MeshOp mesh) {
158  SmallVector<MeshAxis> sorted = llvm::to_vector(axes);
159  llvm::sort(sorted);
160  if (!isUnique(sorted.begin(), sorted.end())) {
161  return emitError(loc) << "Mesh axes contains duplicate elements.";
162  }
163 
164  MeshAxis rank = mesh.getRank();
165  for (auto axis : axes) {
166  if (axis >= rank || axis < 0) {
167  return emitError(loc)
168  << "0-based mesh axis index " << axis
169  << " is out of bounds. The referenced mesh \"" << mesh.getSymName()
170  << "\" is of rank " << rank << ".";
171  }
172  }
173 
174  return success();
175 }
176 
177 template <typename Op>
178 static FailureOr<MeshOp>
180  auto mesh =
181  ::getMeshAndVerify(op.getOperation(), op.getMeshAttr(), symbolTable);
182  if (failed(mesh)) {
183  return failure();
184  }
185  if (failed(verifyMeshAxes(op.getLoc(), op.getMeshAxes(), mesh.value()))) {
186  return failure();
187  }
188  return mesh;
189 }
190 
191 template <typename InShape, typename MeshShape, typename SplitAxes,
192  typename OutShape>
193 static void shardShape(const InShape &inShape, const MeshShape &meshShape,
194  const SplitAxes &splitAxes, OutShape &outShape,
195  ArrayRef<int64_t> shardedDimsOffsets = {},
196  ArrayRef<int64_t> haloSizes = {}) {
197  // 0d tensors cannot be sharded and must get replicated
198  if (inShape.empty()) {
199  assert(outShape.empty());
200  return;
201  }
202 
203  std::copy(llvm::adl_begin(inShape), llvm::adl_end(inShape),
204  llvm::adl_begin(outShape));
205 
206  if (!shardedDimsOffsets.empty()) {
207  auto isDynShape = ShapedType::isDynamicShape(meshShape);
208  uint64_t pos = 1;
209  for (auto [tensorAxis, innerSplitAxes] : llvm::enumerate(splitAxes)) {
210  if (!innerSplitAxes.empty()) {
211  auto sz = shardedDimsOffsets[pos];
212  bool same = !isDynShape;
213  if (same) {
214  // Find sharded dims in shardedDimsOffsets with same static size on
215  // all devices. Use kDynamic for dimensions with dynamic or
216  // non-uniform offs in shardedDimsOffsets.
217  uint64_t numShards = 0;
218  for (auto i : innerSplitAxes.asArrayRef()) {
219  numShards += meshShape[i];
220  }
221  for (size_t i = 1; i < numShards; ++i) {
222  if (shardedDimsOffsets[pos + i] - shardedDimsOffsets[pos + i - 1] !=
223  sz) {
224  same = false;
225  break;
226  }
227  }
228  pos += numShards + 1;
229  }
230  outShape[tensorAxis] = same ? sz : ShapedType::kDynamic;
231  }
232  }
233  } else {
234  for (auto [tensorAxis, innerSplitAxes] : llvm::enumerate(splitAxes)) {
235  outShape[tensorAxis] = shardDimension(
236  inShape[tensorAxis],
237  collectiveProcessGroupSize(innerSplitAxes.asArrayRef(), meshShape));
238  }
239 
240  if (!haloSizes.empty()) {
241  // add halo sizes if requested
242  int haloAxis = 0;
243  for (auto [tensorAxis, innerSplitAxes] : llvm::enumerate(splitAxes)) {
244  if (!ShapedType::isDynamic(outShape[tensorAxis]) &&
245  !innerSplitAxes.empty()) {
246  if (haloSizes[haloAxis * 2] >= 0 &&
247  haloSizes[haloAxis * 2 + 1] >= 0) {
248  outShape[tensorAxis] +=
249  haloSizes[haloAxis * 2] + haloSizes[haloAxis * 2 + 1];
250  ++haloAxis;
251  } else {
252  outShape[tensorAxis] = ShapedType::kDynamic;
253  }
254  }
255  }
256  }
257  }
258 }
259 
260 ShapedType mesh::shardShapedType(ShapedType shape, MeshOp mesh,
261  MeshSharding sharding) {
262  using Dim = std::decay_t<decltype(shape.getDimSize(0))>;
263  SmallVector<Dim> resShapeArr(shape.getShape().size());
264  shardShape(shape.getShape(), mesh.getShape(), sharding.getSplitAxes(),
265  resShapeArr, sharding.getStaticShardedDimsOffsets(),
266  sharding.getStaticHaloSizes());
267  return shape.clone(resShapeArr);
268 }
269 
270 Type mesh::shardType(Type type, MeshOp mesh, MeshSharding sharding) {
271  RankedTensorType rankedTensorType = dyn_cast<RankedTensorType>(type);
272  if (rankedTensorType && !rankedTensorType.getShape().empty()) {
273  return shardShapedType(rankedTensorType, mesh, sharding);
274  }
275  return type;
276 }
277 
279  OpOperand &operand,
280  OpBuilder &builder,
281  ShardOp &newShardOp) {
282  OpBuilder::InsertionGuard insertionGuard(builder);
283  Value operandValue = operand.get();
284  Operation *operandOp = operand.getOwner();
285  builder.setInsertionPointAfterValue(operandValue);
286  ShardOp shardOp = dyn_cast<ShardOp>(operandOp);
287  if (shardOp && sharding == shardOp.getSharding() &&
288  !shardOp.getAnnotateForUsers()) {
289  // No need for anything if the correct sharding is already set.
290  if (!newShardOp) {
291  newShardOp = shardOp;
292  }
293  return;
294  }
295 
296  if (!newShardOp) {
297  auto shardingOp =
298  builder.create<ShardingOp>(operandValue.getLoc(), sharding);
299  newShardOp =
300  builder.create<ShardOp>(operandValue.getLoc(), operandValue, shardingOp,
301  /*annotate_for_users*/ false);
302  }
303  IRRewriter rewriter(builder);
304  rewriter.replaceUsesWithIf(
305  operandValue, newShardOp, [operandOp, operandValue](OpOperand &use) {
306  return use.getOwner() == operandOp && use.get() == operandValue;
307  });
308 
309  if (!shardOp || shardOp.getAnnotateForUsers()) {
310  return;
311  }
312 
313  auto newShardOp2 = builder.create<ShardOp>(operandValue.getLoc(), newShardOp,
314  newShardOp.getSharding(),
315  /*annotate_for_users*/ true);
316  rewriter.replaceAllUsesExcept(newShardOp, newShardOp2, newShardOp2);
317  return;
318 }
319 
321  OpResult result,
322  OpBuilder &builder) {
323  ShardOp newShardOp;
324  for (auto &use : llvm::make_early_inc_range(result.getUses())) {
325  maybeInsertTargetShardingAnnotation(sharding, use, builder, newShardOp);
326  }
327 }
328 
330  OpOperand &operand,
331  OpBuilder &builder) {
332  OpBuilder::InsertionGuard insertionGuard(builder);
333  Value operandValue = operand.get();
334  Operation *operandSrcOp = operandValue.getDefiningOp();
335  bool isBlockArg = !operandSrcOp;
336  {
337  [[maybe_unused]] auto opType =
338  dyn_cast<mlir::RankedTensorType>(operandValue.getType());
339  assert(!opType || opType.getRank() > 0 || isFullReplication(sharding));
340  }
341  if (!isa<RankedTensorType>(operandValue.getType()) && operandSrcOp &&
342  operandSrcOp->hasTrait<OpTrait::ConstantLike>()) {
343  return;
344  }
345 
346  Operation *operandOp = operand.getOwner();
347  ShardOp shardOp = dyn_cast_or_null<ShardOp>(operandSrcOp);
348 
349  if (shardOp && sharding == shardOp.getSharding() &&
350  shardOp.getAnnotateForUsers()) {
351  // No need for anything the correct sharding is already set.
352  return;
353  }
354 
355  builder.setInsertionPoint(operandOp);
356  auto shardingOp =
357  builder.create<ShardingOp>(operand.get().getLoc(), sharding);
358  auto newShardOp =
359  builder.create<ShardOp>(operandValue.getLoc(), operandValue, shardingOp,
360  /*annotate_for_users*/ true);
361  IRRewriter rewriter(builder);
362  rewriter.replaceUsesWithIf(
363  operandValue, newShardOp, [operandOp, operandValue](OpOperand &use) {
364  return use.getOwner() == operandOp && use.get() == operandValue;
365  });
366 
367  if (isBlockArg || !shardOp || !shardOp.getAnnotateForUsers()) {
368  // No need for resharding.
369  return;
370  }
371 
372  builder.setInsertionPoint(newShardOp);
373  auto newPreceedingShardOp =
374  builder.create<ShardOp>(operandValue.getLoc(), operandValue, shardingOp,
375  /*annotate_for_users*/ false);
376  rewriter.replaceUsesWithIf(
377  newShardOp.getSrc(), newPreceedingShardOp, [&newShardOp](OpOperand &use) {
378  return use.getOwner() == newShardOp.getOperation();
379  });
380 }
381 
382 //===----------------------------------------------------------------------===//
383 // mesh.mesh op
384 //===----------------------------------------------------------------------===//
385 
386 LogicalResult MeshOp::verify() {
387  int64_t rank = getRank();
388 
389  if (rank <= 0)
390  return emitOpError("rank of mesh is expected to be a positive integer");
391 
392  for (int64_t dimSize : getShape()) {
393  if (dimSize < 0 && !ShapedType::isDynamic(dimSize))
394  return emitOpError("dimension size of a mesh is expected to be "
395  "non-negative or dynamic");
396  }
397 
398  return success();
399 }
400 
401 //===----------------------------------------------------------------------===//
402 // mesh.mesh_shape op
403 //===----------------------------------------------------------------------===//
404 
405 LogicalResult
406 MeshShapeOp::verifySymbolUses(SymbolTableCollection &symbolTable) {
407  auto mesh = ::getMeshAndVerify(getOperation(), getMeshAttr(), symbolTable);
408  if (failed(mesh)) {
409  return failure();
410  }
411  if (failed(verifyMeshAxes(getLoc(), getAxes(), mesh.value()))) {
412  return failure();
413  }
414 
415  size_t expectedResultsCount =
416  getAxes().empty() ? mesh->getRank() : getAxes().size();
417  if (getResult().size() != expectedResultsCount) {
418  return emitError() << "Unexpected number of results " << getResult().size()
419  << ". Expected " << expectedResultsCount << ".";
420  }
421 
422  return success();
423 }
424 
425 void MeshShapeOp::build(OpBuilder &odsBuilder, OperationState &odsState,
426  MeshOp mesh) {
427  build(odsBuilder, odsState, mesh, SmallVector<MeshAxis>());
428 }
429 
430 void MeshShapeOp::build(OpBuilder &odsBuilder, OperationState &odsState,
431  MeshOp mesh, ArrayRef<MeshAxis> axes) {
432  build(odsBuilder, odsState,
433  SmallVector<Type>(axes.empty() ? mesh.getRank() : axes.size(),
434  odsBuilder.getIndexType()),
435  mesh.getSymName(), MeshAxesAttr::get(odsBuilder.getContext(), axes));
436 }
437 
438 void MeshShapeOp::build(OpBuilder &odsBuilder, OperationState &odsState,
439  StringRef mesh, ArrayRef<MeshAxis> axes) {
440  assert(!axes.empty());
441  build(odsBuilder, odsState,
442  SmallVector<Type>(axes.size(), odsBuilder.getIndexType()), mesh,
443  MeshAxesAttr::get(odsBuilder.getContext(), axes));
444 }
445 
446 void MeshShapeOp::getAsmResultNames(
447  function_ref<void(Value, StringRef)> setNameFn) {
448  setNameFn(getResults()[0], "mesh_shape");
449 }
450 
451 //===----------------------------------------------------------------------===//
452 // mesh.sharding
453 //===----------------------------------------------------------------------===//
454 
455 void ShardingOp::build(::mlir::OpBuilder &b, ::mlir::OperationState &odsState,
456  FlatSymbolRefAttr mesh,
457  ArrayRef<MeshAxesAttr> split_axes,
458  ArrayRef<MeshAxis> partial_axes,
459  mesh::ReductionKind partial_type,
460  ArrayRef<int64_t> static_halos,
461  ArrayRef<int64_t> static_offsets) {
462  return build(
463  b, odsState, mesh, MeshAxesArrayAttr::get(b.getContext(), split_axes),
464  ::mlir::DenseI16ArrayAttr::get(b.getContext(), partial_axes),
465  ::mlir::mesh::ReductionKindAttr::get(b.getContext(), partial_type),
466  ::mlir::DenseI64ArrayAttr::get(b.getContext(), static_halos), {},
467  ::mlir::DenseI64ArrayAttr::get(b.getContext(), static_offsets), {});
468 }
469 
470 void ShardingOp::build(::mlir::OpBuilder &b, ::mlir::OperationState &odsState,
471  FlatSymbolRefAttr mesh,
472  ArrayRef<MeshAxesAttr> split_axes) {
473  return build(
474  b, odsState, mesh, MeshAxesArrayAttr::get(b.getContext(), split_axes), {},
475  ::mlir::mesh::ReductionKindAttr::get(b.getContext(), ReductionKind::Sum),
476  {}, {}, {}, {});
477 }
478 
479 void ShardingOp::build(::mlir::OpBuilder &b, ::mlir::OperationState &odsState,
480  llvm::StringRef mesh, ArrayRef<MeshAxesAttr> split_axes,
481  ArrayRef<int64_t> static_halos,
482  ArrayRef<int64_t> static_offsets) {
483  return build(
484  b, odsState, FlatSymbolRefAttr::get(b.getContext(), mesh),
485  MeshAxesArrayAttr::get(b.getContext(), split_axes), {},
486  ::mlir::mesh::ReductionKindAttr::get(b.getContext(), ReductionKind::Sum),
487  ::mlir::DenseI64ArrayAttr::get(b.getContext(), static_halos), {},
488  ::mlir::DenseI64ArrayAttr::get(b.getContext(), static_offsets), {});
489 }
490 
491 void ShardingOp::build(
492  ::mlir::OpBuilder &b, ::mlir::OperationState &odsState,
493  FlatSymbolRefAttr mesh, ArrayRef<MeshAxesAttr> split_axes,
495  ::mlir::ArrayRef<::mlir::OpFoldResult> sharded_dims_offsets) {
496  mlir::SmallVector<int64_t> staticHalos, staticDims;
497  mlir::SmallVector<mlir::Value> dynamicHalos, dynamicDims;
498  dispatchIndexOpFoldResults(halo_sizes, dynamicHalos, staticHalos);
499  dispatchIndexOpFoldResults(sharded_dims_offsets, dynamicDims, staticDims);
500  return build(
501  b, odsState, mesh, MeshAxesArrayAttr::get(b.getContext(), split_axes), {},
502  ::mlir::mesh::ReductionKindAttr::get(b.getContext(), ReductionKind::Sum),
503  ::mlir::DenseI64ArrayAttr::get(b.getContext(), staticHalos), dynamicHalos,
504  ::mlir::DenseI64ArrayAttr::get(b.getContext(), staticDims), dynamicDims);
505 }
506 
507 void ShardingOp::build(::mlir::OpBuilder &b, ::mlir::OperationState &odsState,
509 
510  build(b, odsState, ShardingType::get(b.getContext()), from.getMeshAttr(),
512  from.getPartialAxes().empty()
516  from.getPartialType()),
517  from.getStaticShardedDimsOffsets().empty()
521  from.getStaticHaloSizes().empty()
524  from.getDynamicHaloSizes());
525 }
526 
527 LogicalResult ShardingOp::verify() {
528  llvm::SmallSet<MeshAxis, 4> visitedAxes;
529 
530  auto checkMeshAxis = [&](ArrayRef<MeshAxis> axesArray) -> LogicalResult {
531  for (MeshAxis axis : axesArray) {
532  if (axis < 0)
533  return emitError() << "mesh axis is expected to be non-negative";
534  if (!visitedAxes.insert(axis).second)
535  return emitError() << "mesh axis duplicated";
536  }
537  return success();
538  };
539 
540  for (auto subAxes : getSplitAxes().getAxes()) {
541  ArrayRef<MeshAxis> subAxesArray = subAxes.asArrayRef();
542  if (failed(checkMeshAxis(subAxesArray)))
543  return failure();
544  }
545  if (getPartialAxes().has_value() &&
546  failed(checkMeshAxis(getPartialAxes().value())))
547  return failure();
548 
549  if (!getStaticHaloSizes().empty() && !getStaticShardedDimsOffsets().empty()) {
550  return emitOpError("halo sizes and shard offsets are mutually exclusive");
551  }
552 
553  if (!getStaticHaloSizes().empty()) {
554  auto numSplitAxes = getSplitAxes().getAxes().size();
555  for (auto splitAxis : getSplitAxes().getAxes()) {
556  if (splitAxis.empty()) {
557  --numSplitAxes;
558  }
559  }
560  if (getStaticHaloSizes().size() != numSplitAxes * 2) {
561  return emitError() << "halo sizes must be specified for all split axes.";
562  }
563  }
564 
565  return success();
566 }
567 
568 void ShardingOp::getAsmResultNames(
569  function_ref<void(Value, StringRef)> setNameFn) {
570  setNameFn(getResult(), "sharding");
571 }
572 
573 LogicalResult ShardingOp::verifySymbolUses(SymbolTableCollection &symbolTable) {
574  auto mesh = ::getMeshAndVerify(getOperation(), getMeshAttr(), symbolTable);
575  if (failed(mesh)) {
576  return failure();
577  }
578  if (mlir::ShapedType::isDynamicShape(mesh->getShape()) &&
579  getStaticShardedDimsOffsets().size() > 0) {
580  return emitError() << "sharded dims offsets are not allowed for "
581  "devices meshes with dynamic shape.";
582  }
583 
584  auto shardedDimsOffsets = getStaticShardedDimsOffsets();
585  if (!shardedDimsOffsets.empty()) {
586  auto meshShape = mesh.value().getShape();
587  assert(!ShapedType::isDynamicShape(meshShape));
588  uint64_t pos = 0;
589  for (auto [tensorAxis, innerSplitAxes] : llvm::enumerate(getSplitAxes())) {
590  if (!innerSplitAxes.empty()) {
591  int64_t numShards = 0, off = 0;
592  for (auto i : innerSplitAxes.asArrayRef()) {
593  numShards += meshShape[i];
594  }
595  for (int64_t i = 0; i <= numShards; ++i) {
596  if (shardedDimsOffsets.size() <= pos + i) {
597  return emitError() << "sharded dims offsets has wrong size.";
598  }
599  if (!ShapedType::isDynamic(shardedDimsOffsets[pos + i])) {
600  if (shardedDimsOffsets[pos + i] < off) {
601  return emitError()
602  << "sharded dims offsets must be non-decreasing.";
603  }
604  off = shardedDimsOffsets[pos + i];
605  }
606  }
607  pos += numShards + 1;
608  }
609  }
610  }
611  return success();
612 }
613 
614 namespace {
615 // Sharding annotations "halo sizes" and "sharded dims offsets"
616 // are a mix of attributes and dynamic values. This canonicalization moves
617 // constant values to the respective attribute lists, minimizing the number
618 // of values.
619 // It also removes sharded_dims_sizes and halos if they are effectively "empty".
620 class NormalizeSharding final : public OpRewritePattern<ShardingOp> {
621 public:
623 
624  LogicalResult matchAndRewrite(ShardingOp op,
625  PatternRewriter &b) const override {
626  auto mixedHalos =
627  getMixedValues(op.getStaticHaloSizes(), op.getDynamicHaloSizes(), b);
628  auto mixedOffs = getMixedValues(op.getStaticShardedDimsOffsets(),
629  op.getDynamicShardedDimsOffsets(), b);
630 
631  // No constant operands were folded, just return;
632  bool modified = succeeded(foldDynamicIndexList(mixedHalos, true)) ||
633  succeeded(foldDynamicIndexList(mixedOffs, true));
634 
635  auto [staticHalos, dynamicHalos] = decomposeMixedValues(mixedHalos);
636  auto [staticOffs, dynamicOffs] = decomposeMixedValues(mixedOffs);
637 
638  if (dynamicHalos.empty() && !staticHalos.empty()) {
639  if (staticHalos[0] == 0 && llvm::all_equal(staticHalos)) {
640  staticHalos.clear();
641  modified = true;
642  }
643  }
644 
645  // Remove sharded dims offsets if they are effectively the default values,
646  // e.g. if they define equi-distance between all neighboring shards.
647  // Requires static-only offsets. Compares the first distance as the
648  // difference between the first two offsets. Only if all consecutive
649  // distances are the same, the offsets are removed.
650  if (dynamicOffs.empty() && !staticOffs.empty()) {
651  assert(staticOffs.size() >= 2);
652  auto diff = staticOffs[1] - staticOffs[0];
653  bool all_same = staticOffs.size() > 2;
654  for (auto i = 2u; i < staticOffs.size(); ++i) {
655  if (staticOffs[i] - staticOffs[i - 1] != diff) {
656  all_same = false;
657  break;
658  }
659  }
660  if (all_same) {
661  staticOffs.clear();
662  modified = true;
663  }
664  }
665 
666  if (!modified) {
667  return failure();
668  }
669 
670  op.setStaticHaloSizes(staticHalos);
671  op.getDynamicHaloSizesMutable().assign(dynamicHalos);
672  op.setStaticShardedDimsOffsets(staticOffs);
673  op.getDynamicShardedDimsOffsetsMutable().assign(dynamicOffs);
674 
675  return success();
676  }
677 };
678 } // namespace
679 
680 void ShardingOp::getCanonicalizationPatterns(mlir::RewritePatternSet &results,
681  mlir::MLIRContext *context) {
682  results.add<NormalizeSharding>(context);
683 }
684 
685 //===----------------------------------------------------------------------===//
686 // MeshSharding
687 //===----------------------------------------------------------------------===//
688 
690  if (getMesh() != rhs.getMesh()) {
691  return false;
692  }
693 
694  if (getPartialAxes().size() != rhs.getPartialAxes().size() ||
695  (!getPartialAxes().empty() && getPartialType() != rhs.getPartialType()) ||
696  !llvm::equal(getPartialAxes(), rhs.getPartialAxes())) {
697  return false;
698  }
699 
700  auto minSize = std::min(getSplitAxes().size(), rhs.getSplitAxes().size());
701  if (!llvm::equal(llvm::make_range(getSplitAxes().begin(),
702  getSplitAxes().begin() + minSize),
703  llvm::make_range(rhs.getSplitAxes().begin(),
704  rhs.getSplitAxes().begin() + minSize))) {
705  return false;
706  }
707 
708  return llvm::all_of(llvm::drop_begin(getSplitAxes(), minSize),
709  std::mem_fn(&MeshAxesAttr::empty)) &&
710  llvm::all_of(llvm::drop_begin(rhs.getSplitAxes(), minSize),
711  std::mem_fn(&MeshAxesAttr::empty));
712 }
713 
715  return equalShardSizes(rhs) && equalHaloSizes(rhs);
716 }
717 
719  if (rhs.getStaticShardedDimsOffsets().size() !=
720  getStaticShardedDimsOffsets().size() ||
721  !llvm::equal(getStaticShardedDimsOffsets(),
723  return false;
724  }
725  if (rhs.getDynamicShardedDimsOffsets().size() !=
726  getDynamicShardedDimsOffsets().size() ||
727  !llvm::equal(getDynamicShardedDimsOffsets(),
729  return false;
730  }
731  return true;
732 }
733 
735  if (rhs.getStaticHaloSizes().size() != getStaticHaloSizes().size() ||
736  !llvm::equal(getStaticHaloSizes(), rhs.getStaticHaloSizes())) {
737  return false;
738  }
739  if (rhs.getDynamicHaloSizes().size() != getDynamicHaloSizes().size() ||
740  !llvm::equal(getDynamicHaloSizes(), rhs.getDynamicHaloSizes())) {
741  return false;
742  }
743  return true;
744 }
745 
748 }
749 
750 bool MeshSharding::operator!=(Value rhs) const { return !(*this == rhs); }
751 
752 bool MeshSharding::operator==(const MeshSharding &rhs) const {
754 }
755 
756 bool MeshSharding::operator!=(const MeshSharding &rhs) const {
757  return !(*this == rhs);
758 }
759 
761 
763  auto shardingOp = mlir::dyn_cast<ShardingOp>(rhs.getDefiningOp());
764  assert(shardingOp && "expected sharding op");
765  auto splitAxes = shardingOp.getSplitAxes().getAxes();
766  auto partialAxes = shardingOp.getPartialAxes().value_or(ArrayRef<MeshAxis>());
767  // If splitAxes and partialAxes are empty, use "empty" constructor.
768  if (splitAxes.empty() && partialAxes.empty()) {
769  *this = MeshSharding(shardingOp.getMeshAttr());
770  return;
771  }
772  *this = get(shardingOp.getMeshAttr(), splitAxes, partialAxes,
773  shardingOp.getPartialType().value_or(ReductionKind::Sum),
774  shardingOp.getStaticHaloSizes(),
775  shardingOp.getStaticShardedDimsOffsets(),
776  SmallVector<Value>(shardingOp.getDynamicHaloSizes()),
777  SmallVector<Value>(shardingOp.getDynamicShardedDimsOffsets()));
778 }
779 
781  ArrayRef<MeshAxesAttr> split_axes_,
782  ArrayRef<MeshAxis> partial_axes_,
783  ReductionKind partial_type_,
784  ArrayRef<int64_t> static_halo_sizes_,
785  ArrayRef<int64_t> static_sharded_dims_offsets_,
786  ArrayRef<Value> dynamic_halo_sizes_,
787  ArrayRef<Value> dynamic_sharded_dims_offsets_) {
788  MeshSharding res(mesh_);
789  if (split_axes_.empty() && partial_axes_.empty()) {
790  return res;
791  }
792 
793  res.split_axes.resize(split_axes_.size());
794  for (auto [i, axis] : llvm::enumerate(split_axes_)) {
795  res.split_axes[i] =
796  MeshAxesAttr::get(mesh_.getContext(), axis.asArrayRef());
797  }
798 
799  auto clone = [](const auto src, auto &dst) {
800  dst.resize(src.size());
801  llvm::copy(src, dst.begin());
802  };
803 
804  clone(partial_axes_, res.partial_axes);
805  res.partial_type = partial_type_;
806  clone(static_halo_sizes_, res.static_halo_sizes);
807  clone(static_sharded_dims_offsets_, res.static_sharded_dims_offsets);
808  clone(dynamic_halo_sizes_, res.dynamic_halo_sizes);
809  clone(dynamic_sharded_dims_offsets_, res.dynamic_sharded_dims_offsets);
810 
811  return res;
812 }
813 
814 //===----------------------------------------------------------------------===//
815 // mesh.shard_shape
816 //===----------------------------------------------------------------------===//
817 
818 void ShardShapeOp::getAsmResultNames(
819  function_ref<void(Value, StringRef)> setNameFn) {
820  setNameFn(getResult()[0], "shard_shape");
821 }
822 
823 void ShardShapeOp::build(::mlir::OpBuilder &odsBuilder,
824  ::mlir::OperationState &odsState,
826  ArrayRef<Value> dims_dyn, ::mlir::Value sharding,
827  ::mlir::ValueRange device) {
828  SmallVector<mlir::Type> resType(dims.size(), odsBuilder.getIndexType());
829  build(odsBuilder, odsState, resType, dims, dims_dyn, sharding,
830  SmallVector<int64_t>(device.size(), ShapedType::kDynamic), device);
831 }
832 
833 //===----------------------------------------------------------------------===//
834 // mesh.shard op
835 //===----------------------------------------------------------------------===//
836 
837 void ShardOp::getAsmResultNames(
838  function_ref<void(Value, StringRef)> setNameFn) {
839  setNameFn(getResult(), "sharding_annotated");
840 }
841 
842 namespace {
843 // Determine if the given ShardOp is a duplicate of another ShardOp
844 // on the same value. This can happen if constant values are sharded.
845 class FoldDuplicateShardOp final : public OpRewritePattern<ShardOp> {
846 public:
848 
849  LogicalResult matchAndRewrite(ShardOp op, PatternRewriter &b) const override {
850  // Get the use-list of the value being sharded and check if it has more than
851  // one use.
852  Value value = op.getSrc();
853  if (value.hasOneUse() || value.getDefiningOp<ShardOp>()) {
854  return failure();
855  }
856 
857  // Iterate through the uses of the value to find a duplicate ShardOp.
858  for (auto &use : value.getUses()) {
859  if (use.getOwner() != op.getOperation()) {
860  auto otherOp = dyn_cast<ShardOp>(use.getOwner());
861  if (!otherOp || !otherOp->isBeforeInBlock(op)) {
862  return failure();
863  }
864  // Create a MeshSharding object for the current and the other ShardOp
865  // If the two are equal replace current op with the other op.
866  MeshSharding currentSharding(op.getSharding());
867  MeshSharding otherSharding(otherOp.getSharding());
868  if (currentSharding == otherSharding) {
869  b.replaceAllUsesWith(op.getResult(), otherOp.getResult());
870  b.eraseOp(op.getOperation());
871  } else {
872  // use the other sharding as input for op
873  op.getSrcMutable().assign(otherOp.getResult());
874  }
875  return success();
876  }
877  }
878 
879  return failure();
880  }
881 };
882 } // namespace
883 
884 void ShardOp::getCanonicalizationPatterns(mlir::RewritePatternSet &results,
885  mlir::MLIRContext *context) {
886  results.add<FoldDuplicateShardOp>(context);
887 }
888 
889 //===----------------------------------------------------------------------===//
890 // mesh.process_multi_index op
891 //===----------------------------------------------------------------------===//
892 
893 LogicalResult
894 ProcessMultiIndexOp::verifySymbolUses(SymbolTableCollection &symbolTable) {
895  auto mesh = ::getMeshAndVerify(getOperation(), getMeshAttr(), symbolTable);
896  if (failed(mesh)) {
897  return failure();
898  }
899  if (failed(verifyMeshAxes(getLoc(), getAxes(), mesh.value()))) {
900  return failure();
901  }
902 
903  size_t expectedResultsCount =
904  getAxes().empty() ? mesh->getRank() : getAxes().size();
905  if (getResult().size() != expectedResultsCount) {
906  return emitError() << "Unexpected number of results " << getResult().size()
907  << ". Expected " << expectedResultsCount << ".";
908  }
909 
910  return success();
911 }
912 
913 void ProcessMultiIndexOp::build(OpBuilder &odsBuilder, OperationState &odsState,
914  MeshOp mesh) {
915  build(odsBuilder, odsState,
916  SmallVector<Type>(mesh.getRank(), odsBuilder.getIndexType()),
917  mesh.getSymName(), ArrayRef<MeshAxis>());
918 }
919 
920 void ProcessMultiIndexOp::build(OpBuilder &odsBuilder, OperationState &odsState,
921  StringRef mesh, ArrayRef<MeshAxis> axes) {
922  build(odsBuilder, odsState,
923  SmallVector<Type>(axes.size(), odsBuilder.getIndexType()), mesh,
924  MeshAxesAttr::get(odsBuilder.getContext(), axes));
925 }
926 
927 void ProcessMultiIndexOp::getAsmResultNames(
928  function_ref<void(Value, StringRef)> setNameFn) {
929  setNameFn(getResults()[0], "proc_linear_idx");
930 }
931 
932 //===----------------------------------------------------------------------===//
933 // mesh.process_linear_index op
934 //===----------------------------------------------------------------------===//
935 
936 LogicalResult
937 ProcessLinearIndexOp::verifySymbolUses(SymbolTableCollection &symbolTable) {
938  auto mesh = ::getMeshAndVerify(getOperation(), getMeshAttr(), symbolTable);
939  if (failed(mesh)) {
940  return failure();
941  }
942  return success();
943 }
944 
945 void ProcessLinearIndexOp::build(OpBuilder &odsBuilder,
946  OperationState &odsState, MeshOp mesh) {
947  build(odsBuilder, odsState, mesh.getSymName());
948 }
949 
950 void ProcessLinearIndexOp::getAsmResultNames(
951  function_ref<void(Value, StringRef)> setNameFn) {
952  setNameFn(getResult(), "proc_linear_idx");
953 }
954 
955 //===----------------------------------------------------------------------===//
956 // mesh.neighbors_linear_indices op
957 //===----------------------------------------------------------------------===//
958 
959 LogicalResult
960 NeighborsLinearIndicesOp::verifySymbolUses(SymbolTableCollection &symbolTable) {
961  auto mesh = ::getMeshAndVerify(getOperation(), getMeshAttr(), symbolTable);
962  if (failed(mesh)) {
963  return failure();
964  }
965  return success();
966 }
967 
968 void NeighborsLinearIndicesOp::getAsmResultNames(
969  function_ref<void(Value, StringRef)> setNameFn) {
970  setNameFn(getNeighborDown(), "down_linear_idx");
971  setNameFn(getNeighborUp(), "up_linear_idx");
972 }
973 
974 //===----------------------------------------------------------------------===//
975 // collective communication ops
976 //===----------------------------------------------------------------------===//
977 
978 namespace {
979 
980 template <typename Op>
981 struct EmptyMeshAxesCanonicalizationPattern : OpRewritePattern<Op> {
983  LogicalResult matchAndRewrite(Op op,
984  PatternRewriter &rewriter) const override {
985  auto meshAxes = op.getMeshAxes();
986  if (!meshAxes.empty()) {
987  return failure();
988  }
989  if (op.getInput().getType() != op.getResult().getType()) {
990  return failure();
991  }
992 
993  rewriter.replaceAllUsesWith(op.getResult(), op.getInput());
994  rewriter.eraseOp(op.getOperation());
995  return success();
996  }
997 };
998 
999 } // namespace
1000 
1001 static LogicalResult verifyInGroupDevice(Location loc, StringRef deviceName,
1002  ArrayRef<int64_t> device,
1003  Operation::operand_range deviceDynamic,
1004  ArrayRef<MeshAxis> meshAxes,
1005  ArrayRef<int64_t> meshShape) {
1006  if (device.size() != meshAxes.size()) {
1007  return emitError(loc) << "In-group device \"" << deviceName
1008  << "\" has unexpected multi-index size "
1009  << device.size() << ". Expected " << meshAxes.size()
1010  << ".";
1011  }
1012 
1013  for (size_t i = 0; i < device.size(); ++i) {
1014  if (!ShapedType::isDynamic(device[i]) &&
1015  !ShapedType::isDynamic(meshShape[meshAxes[i]]) &&
1016  meshShape[meshAxes[i]] <= device[i]) {
1017  return emitError(loc)
1018  << "Out of bounds coordinate " << i << " for in-group device \""
1019  << deviceName << "\"."
1020  << " Got " << device[i] << ", but expected value in the range [0, "
1021  << (meshShape[meshAxes[i]] - 1) << "].";
1022  }
1023  }
1024  return success();
1025 }
1026 
1027 template <typename It>
1028 static auto product(It begin, It end) {
1029  using ElementType = std::decay_t<decltype(*begin)>;
1030  return std::accumulate(begin, end, static_cast<ElementType>(1),
1031  std::multiplies<ElementType>());
1032 }
1033 
1034 template <typename R>
1035 static auto product(R &&range) {
1036  return product(adl_begin(range), adl_end(range));
1037 }
1038 
1039 static LogicalResult verifyDimensionCompatibility(Location loc,
1040  int64_t expectedDimSize,
1041  int64_t resultDimSize,
1042  int64_t resultAxis) {
1043  if (!ShapedType::isDynamic(resultDimSize) &&
1044  expectedDimSize != resultDimSize) {
1045  return emitError(loc) << "Dimension size mismatch for result axis "
1046  << resultAxis << ". Expected "
1047  << (ShapedType::isDynamic(expectedDimSize)
1048  ? Twine("dynamic")
1049  : Twine(expectedDimSize))
1050  << ", but got " << resultDimSize << ".";
1051  }
1052 
1053  return success();
1054 }
1055 
1057  Value operand, Value result, int64_t gatherAxis,
1058  ArrayRef<MeshAxis> meshAxes, ArrayRef<int64_t> meshShape) {
1059  auto resultRank = cast<ShapedType>(result.getType()).getRank();
1060  if (gatherAxis < 0 || gatherAxis >= resultRank) {
1061  return emitError(result.getLoc())
1062  << "Gather axis " << gatherAxis << " is out of bounds [0, "
1063  << resultRank << ").";
1064  }
1065 
1066  ShapedType operandType = cast<ShapedType>(operand.getType());
1067  ShapedType resultType = cast<ShapedType>(result.getType());
1068  auto deviceGroupSize =
1069  DimensionSize(collectiveProcessGroupSize(meshAxes, meshShape));
1070  for (int64_t axis = 0; axis < operandType.getRank(); ++axis) {
1071  auto operandDimSize = DimensionSize(operandType.getDimSize(axis));
1072  auto resultDimSize = DimensionSize(resultType.getDimSize(axis));
1073  auto expectedResultDimSize =
1074  axis == gatherAxis ? deviceGroupSize * operandDimSize : operandDimSize;
1075  if (failed(verifyDimensionCompatibility(
1076  result.getLoc(), expectedResultDimSize, resultDimSize, axis))) {
1077  return failure();
1078  }
1079  }
1080  return success();
1081 }
1082 
1084  Value operand, Value result, int64_t splitAxis, int64_t concatAxis,
1085  ArrayRef<MeshAxis> meshAxes, ArrayRef<int64_t> meshShape) {
1086  ShapedType operandType = cast<ShapedType>(operand.getType());
1087  ShapedType resultType = cast<ShapedType>(result.getType());
1088  for (int64_t axis = 0; axis < operandType.getRank(); ++axis) {
1089  if ((axis != splitAxis && axis != concatAxis) || splitAxis == concatAxis) {
1090  if (failed(verifyDimensionCompatibility(
1091  result.getLoc(), operandType.getDimSize(axis),
1092  resultType.getDimSize(axis), axis))) {
1093  return failure();
1094  }
1095  }
1096  }
1097 
1098  if (splitAxis == concatAxis) {
1099  return success();
1100  }
1101 
1102  auto deviceGroupSize =
1103  DimensionSize(collectiveProcessGroupSize(meshAxes, meshShape));
1104  auto operandConcatDimSize = DimensionSize(operandType.getDimSize(concatAxis));
1105  auto operandSplitDimSize = DimensionSize(operandType.getDimSize(splitAxis));
1106  DimensionSize expectedResultConcatDimSize =
1107  operandConcatDimSize * deviceGroupSize;
1108  DimensionSize expectedResultSplitDimSize =
1109  operandSplitDimSize / deviceGroupSize;
1110  if (!expectedResultSplitDimSize.isDynamic() &&
1111  int64_t(operandSplitDimSize) % int64_t(deviceGroupSize) != 0) {
1112  expectedResultSplitDimSize = DimensionSize::dynamic();
1113  }
1114  if (failed(verifyDimensionCompatibility(
1115  result.getLoc(), expectedResultConcatDimSize.value(),
1116  resultType.getDimSize(concatAxis), concatAxis))) {
1117  return failure();
1118  }
1119  if (failed(verifyDimensionCompatibility(
1120  result.getLoc(), expectedResultSplitDimSize.value(),
1121  resultType.getDimSize(splitAxis), splitAxis))) {
1122  return failure();
1123  }
1124 
1125  return success();
1126 }
1127 
1129  Value operand, Value result, int64_t tensorAxis,
1130  ArrayRef<MeshAxis> meshAxes, ArrayRef<int64_t> meshShape) {
1131  ShapedType operandType = cast<ShapedType>(operand.getType());
1132  ShapedType resultType = cast<ShapedType>(result.getType());
1133  for (int64_t axis = 0; axis < operandType.getRank(); ++axis) {
1134  if (axis != tensorAxis) {
1135  if (failed(verifyDimensionCompatibility(
1136  result.getLoc(), operandType.getDimSize(axis),
1137  resultType.getDimSize(axis), axis))) {
1138  return failure();
1139  }
1140  }
1141  }
1142 
1143  auto deviceGroupSize =
1144  DimensionSize(collectiveProcessGroupSize(meshAxes, meshShape));
1145  auto operandScatterDimSize =
1146  DimensionSize(operandType.getDimSize(tensorAxis));
1147  if (!operandScatterDimSize.isDynamic() && !deviceGroupSize.isDynamic() &&
1148  int64_t(operandScatterDimSize) % int64_t(deviceGroupSize) != 0) {
1149  return emitError(result.getLoc())
1150  << "Operand dimension size " << int64_t(operandScatterDimSize)
1151  << " is not divisible by collective device group size "
1152  << int64_t(deviceGroupSize) << " for tensor axis " << tensorAxis
1153  << ".";
1154  }
1155  DimensionSize expectedResultTensorDimSize =
1156  operandScatterDimSize / deviceGroupSize;
1157  if (failed(verifyDimensionCompatibility(
1158  result.getLoc(), expectedResultTensorDimSize.value(),
1159  resultType.getDimSize(tensorAxis), tensorAxis))) {
1160  return failure();
1161  }
1162 
1163  return success();
1164 }
1165 
1166 static RankedTensorType sliceResultType(Type operandType, MeshOp mesh,
1167  ArrayRef<MeshAxis> meshAxes,
1168  int64_t sliceAxis) {
1169  RankedTensorType operandRankedTensorType =
1170  cast<RankedTensorType>(operandType);
1171  DimensionSize operandSliceAxisSize =
1172  operandRankedTensorType.getShape()[sliceAxis];
1173  SmallVector<int64_t> resultShape =
1174  llvm::to_vector(operandRankedTensorType.getShape());
1175 
1176  resultShape[sliceAxis] =
1177  operandSliceAxisSize /
1178  DimensionSize(collectiveProcessGroupSize(meshAxes, mesh));
1179  return operandRankedTensorType.clone(resultShape);
1180 }
1181 
1182 //===----------------------------------------------------------------------===//
1183 // mesh.all_gather op
1184 //===----------------------------------------------------------------------===//
1185 
1186 LogicalResult
1187 AllGatherOp::verifySymbolUses(SymbolTableCollection &symbolTable) {
1188  auto mesh = getMeshAndVerifyAxes(*this, symbolTable);
1189  if (failed(mesh)) {
1190  return failure();
1191  }
1192  auto gatherAxis = getGatherAxis().getSExtValue();
1193  return verifyGatherOperandAndResultShape(getOperand(), getResult(),
1194  gatherAxis, getMeshAxes(),
1195  mesh.value().getShape());
1196 }
1197 
1198 void AllGatherOp::getCanonicalizationPatterns(RewritePatternSet &patterns,
1199  MLIRContext *context) {
1200  patterns.add<EmptyMeshAxesCanonicalizationPattern<AllGatherOp>>(context);
1201 }
1202 
1203 void AllGatherOp::getAsmResultNames(
1204  function_ref<void(Value, StringRef)> setNameFn) {
1205  setNameFn(getResult(), "all_gather");
1206 }
1207 
1208 //===----------------------------------------------------------------------===//
1209 // mesh.all_reduce op
1210 //===----------------------------------------------------------------------===//
1211 
1212 LogicalResult
1213 AllReduceOp::verifySymbolUses(SymbolTableCollection &symbolTable) {
1214  return getMeshAndVerifyAxes(*this, symbolTable);
1215 }
1216 
1217 void AllReduceOp::getCanonicalizationPatterns(RewritePatternSet &patterns,
1218  MLIRContext *context) {
1219  patterns.add<EmptyMeshAxesCanonicalizationPattern<AllReduceOp>>(context);
1220 }
1221 
1222 void AllReduceOp::build(OpBuilder &odsBuilder, OperationState &odsState,
1223  Value input, StringRef mesh,
1224  ArrayRef<MeshAxis> meshAxes, ReductionKind reduction) {
1225  build(odsBuilder, odsState, input.getType(), mesh, meshAxes, input,
1226  reduction);
1227 }
1228 
1229 void AllReduceOp::getAsmResultNames(
1230  function_ref<void(Value, StringRef)> setNameFn) {
1231  setNameFn(getResult(), "all_reduce");
1232 }
1233 
1234 //===----------------------------------------------------------------------===//
1235 // mesh.all_slice op
1236 //===----------------------------------------------------------------------===//
1237 
1238 LogicalResult AllSliceOp::verifySymbolUses(SymbolTableCollection &symbolTable) {
1239  auto mesh = getMeshAndVerifyAxes(*this, symbolTable);
1240  if (failed(mesh)) {
1241  return failure();
1242  }
1244  getOperand(), getResult(), getSliceAxis().getSExtValue(), getMeshAxes(),
1245  mesh.value().getShape());
1246 }
1247 
1248 void AllSliceOp::getCanonicalizationPatterns(RewritePatternSet &patterns,
1249  MLIRContext *context) {
1250  patterns.add<EmptyMeshAxesCanonicalizationPattern<AllSliceOp>>(context);
1251 }
1252 
1253 void AllSliceOp::build(OpBuilder &odsBuilder, OperationState &odsState,
1254  Value input, MeshOp mesh, ArrayRef<MeshAxis> meshAxes,
1255  int64_t sliceAxis) {
1256  Type resultType = sliceResultType(input.getType(), mesh, meshAxes, sliceAxis);
1257  build(odsBuilder, odsState, resultType, input, mesh.getSymName(), meshAxes,
1258  sliceAxis);
1259 }
1260 
1261 void AllSliceOp::build(OpBuilder &odsBuilder, OperationState &odsState,
1262  Type resultType, Value input, StringRef mesh,
1263  ArrayRef<MeshAxis> meshAxes, int64_t sliceAxis) {
1264  build(odsBuilder, odsState, resultType, mesh, meshAxes, input,
1265  APInt(sizeof(sliceAxis) * CHAR_BIT, sliceAxis));
1266 }
1267 
1268 void AllSliceOp::getAsmResultNames(
1269  function_ref<void(Value, StringRef)> setNameFn) {
1270  setNameFn(getResult(), "all_slice");
1271 }
1272 
1273 //===----------------------------------------------------------------------===//
1274 // mesh.all_to_all op
1275 //===----------------------------------------------------------------------===//
1276 
1277 LogicalResult AllToAllOp::verifySymbolUses(SymbolTableCollection &symbolTable) {
1278  auto mesh = getMeshAndVerifyAxes(*this, symbolTable);
1279  if (failed(mesh)) {
1280  return failure();
1281  }
1282 
1284  getOperand(), getResult(), getSplitAxis().getSExtValue(),
1285  getConcatAxis().getSExtValue(), getMeshAxes(), mesh.value().getShape());
1286 }
1287 
1288 void AllToAllOp::getCanonicalizationPatterns(RewritePatternSet &patterns,
1289  MLIRContext *context) {
1290  patterns.add<EmptyMeshAxesCanonicalizationPattern<AllToAllOp>>(context);
1291 }
1292 
1293 void AllToAllOp::getAsmResultNames(
1294  function_ref<void(Value, StringRef)> setNameFn) {
1295  setNameFn(getResult(), "all_to_all");
1296 }
1297 
1298 //===----------------------------------------------------------------------===//
1299 // mesh.broadcast op
1300 //===----------------------------------------------------------------------===//
1301 
1302 LogicalResult
1303 BroadcastOp::verifySymbolUses(SymbolTableCollection &symbolTable) {
1304  auto mesh = getMeshAndVerifyAxes(*this, symbolTable);
1305  if (failed(mesh)) {
1306  return failure();
1307  }
1308  if (failed(verifyInGroupDevice(getLoc(), getRootAttrName(), getRoot(),
1309  getRootDynamic(), getMeshAxes(),
1310  mesh.value().getShape()))) {
1311  return failure();
1312  }
1313 
1314  return success();
1315 }
1316 
1317 void BroadcastOp::getCanonicalizationPatterns(RewritePatternSet &patterns,
1318  MLIRContext *context) {
1319  patterns.add<EmptyMeshAxesCanonicalizationPattern<BroadcastOp>>(context);
1320 }
1321 
1322 void BroadcastOp::getAsmResultNames(
1323  function_ref<void(Value, StringRef)> setNameFn) {
1324  setNameFn(getResult(), "broadcast");
1325 }
1326 
1327 //===----------------------------------------------------------------------===//
1328 // mesh.gather op
1329 //===----------------------------------------------------------------------===//
1330 
1331 LogicalResult GatherOp::verifySymbolUses(SymbolTableCollection &symbolTable) {
1332  auto mesh = getMeshAndVerifyAxes(*this, symbolTable);
1333  if (failed(mesh)) {
1334  return failure();
1335  }
1336  if (failed(verifyInGroupDevice(getLoc(), getRootAttrName(), getRoot(),
1337  getRootDynamic(), getMeshAxes(),
1338  mesh.value().getShape()))) {
1339  return failure();
1340  }
1341 
1342  auto gatherAxis = getGatherAxis().getSExtValue();
1343  return verifyGatherOperandAndResultShape(getInput(), getResult(), gatherAxis,
1344  getMeshAxes(),
1345  mesh.value().getShape());
1346 }
1347 
1348 void GatherOp::getCanonicalizationPatterns(RewritePatternSet &patterns,
1349  MLIRContext *context) {
1350  patterns.add<EmptyMeshAxesCanonicalizationPattern<GatherOp>>(context);
1351 }
1352 
1353 void GatherOp::getAsmResultNames(
1354  function_ref<void(Value, StringRef)> setNameFn) {
1355  setNameFn(getResult(), "gather");
1356 }
1357 
1358 //===----------------------------------------------------------------------===//
1359 // mesh.recv op
1360 //===----------------------------------------------------------------------===//
1361 
1362 LogicalResult RecvOp::verifySymbolUses(SymbolTableCollection &symbolTable) {
1363  auto mesh = getMeshAndVerifyAxes(*this, symbolTable);
1364  if (failed(mesh)) {
1365  return failure();
1366  }
1367  if (getSource() &&
1368  failed(verifyInGroupDevice(getLoc(), getSourceAttrName(),
1369  getSource().value(), getSourceDynamic(),
1370  getMeshAxes(), mesh.value().getShape()))) {
1371  return failure();
1372  }
1373  return success();
1374 }
1375 
1376 void RecvOp::getCanonicalizationPatterns(RewritePatternSet &patterns,
1377  MLIRContext *context) {
1378  patterns.add<EmptyMeshAxesCanonicalizationPattern<RecvOp>>(context);
1379 }
1380 
1381 void RecvOp::getAsmResultNames(function_ref<void(Value, StringRef)> setNameFn) {
1382  setNameFn(getResult(), "recv");
1383 }
1384 
1385 //===----------------------------------------------------------------------===//
1386 // mesh.reduce op
1387 //===----------------------------------------------------------------------===//
1388 
1389 LogicalResult ReduceOp::verifySymbolUses(SymbolTableCollection &symbolTable) {
1390  auto mesh = getMeshAndVerifyAxes(*this, symbolTable);
1391  if (failed(mesh)) {
1392  return failure();
1393  }
1394  if (failed(verifyInGroupDevice(getLoc(), getRootAttrName(), getRoot(),
1395  getRootDynamic(), getMeshAxes(),
1396  mesh.value().getShape()))) {
1397  return failure();
1398  }
1399 
1400  return success();
1401 }
1402 
1403 void ReduceOp::getCanonicalizationPatterns(RewritePatternSet &patterns,
1404  MLIRContext *context) {
1405  patterns.add<EmptyMeshAxesCanonicalizationPattern<ReduceOp>>(context);
1406 }
1407 
1408 void ReduceOp::getAsmResultNames(
1409  function_ref<void(Value, StringRef)> setNameFn) {
1410  setNameFn(getResult(), "reduce");
1411 }
1412 
1413 //===----------------------------------------------------------------------===//
1414 // mesh.reduce_scatter op
1415 //===----------------------------------------------------------------------===//
1416 
1417 LogicalResult
1418 ReduceScatterOp::verifySymbolUses(SymbolTableCollection &symbolTable) {
1419  auto mesh = getMeshAndVerifyAxes(*this, symbolTable);
1420  if (failed(mesh)) {
1421  return failure();
1422  }
1423 
1425  getOperand(), getResult(), getScatterAxis().getSExtValue(), getMeshAxes(),
1426  mesh.value().getShape());
1427 }
1428 
1429 void ReduceScatterOp::getCanonicalizationPatterns(RewritePatternSet &patterns,
1430  MLIRContext *context) {
1431  patterns.add<EmptyMeshAxesCanonicalizationPattern<ReduceScatterOp>>(context);
1432 }
1433 
1434 void ReduceScatterOp::getAsmResultNames(
1435  function_ref<void(Value, StringRef)> setNameFn) {
1436  setNameFn(getResult(), "reduce_scatter");
1437 }
1438 
1439 //===----------------------------------------------------------------------===//
1440 // mesh.scatter op
1441 //===----------------------------------------------------------------------===//
1442 
1443 LogicalResult ScatterOp::verifySymbolUses(SymbolTableCollection &symbolTable) {
1444  auto mesh = getMeshAndVerifyAxes(*this, symbolTable);
1445  if (failed(mesh)) {
1446  return failure();
1447  }
1448  if (failed(verifyInGroupDevice(getLoc(), getRootAttrName(), getRoot(),
1449  getRootDynamic(), getMeshAxes(),
1450  mesh.value().getShape()))) {
1451  return failure();
1452  }
1453 
1454  auto scatterAxis = getScatterAxis().getSExtValue();
1455  return verifyScatterOrSliceOperandAndResultShape(getInput(), getResult(),
1456  scatterAxis, getMeshAxes(),
1457  mesh.value().getShape());
1458 }
1459 
1460 void ScatterOp::getCanonicalizationPatterns(RewritePatternSet &patterns,
1461  MLIRContext *context) {
1462  patterns.add<EmptyMeshAxesCanonicalizationPattern<ScatterOp>>(context);
1463 }
1464 
1465 void ScatterOp::getAsmResultNames(
1466  function_ref<void(Value, StringRef)> setNameFn) {
1467  setNameFn(getResult(), "scatter");
1468 }
1469 
1470 //===----------------------------------------------------------------------===//
1471 // mesh.send op
1472 //===----------------------------------------------------------------------===//
1473 
1474 LogicalResult SendOp::verifySymbolUses(SymbolTableCollection &symbolTable) {
1475  auto mesh = getMeshAndVerifyAxes(*this, symbolTable);
1476  if (failed(mesh)) {
1477  return failure();
1478  }
1479  if (failed(verifyInGroupDevice(getLoc(), getDestinationAttrName(),
1480  getDestination(), getDestinationDynamic(),
1481  getMeshAxes(), mesh.value().getShape()))) {
1482  return failure();
1483  }
1484  return success();
1485 }
1486 
1487 void SendOp::getCanonicalizationPatterns(RewritePatternSet &patterns,
1488  MLIRContext *context) {
1489  patterns.add<EmptyMeshAxesCanonicalizationPattern<SendOp>>(context);
1490 }
1491 
1492 void SendOp::getAsmResultNames(function_ref<void(Value, StringRef)> setNameFn) {
1493  setNameFn(getResult(), "send");
1494 }
1495 
1496 //===----------------------------------------------------------------------===//
1497 // mesh.shift op
1498 //===----------------------------------------------------------------------===//
1499 
1500 LogicalResult ShiftOp::verifySymbolUses(SymbolTableCollection &symbolTable) {
1501  auto mesh = getMeshAndVerifyAxes(*this, symbolTable);
1502  if (failed(mesh)) {
1503  return failure();
1504  }
1505 
1506  auto meshAxes = getMeshAxes();
1507  auto shiftAxis = getShiftAxis().getZExtValue();
1508  if (llvm::find(meshAxes, shiftAxis) == meshAxes.end()) {
1509  return emitError() << "Invalid shift axis " << shiftAxis
1510  << ". It must be one of the grouping mesh axes.";
1511  }
1512 
1513  return success();
1514 }
1515 
1516 void ShiftOp::getCanonicalizationPatterns(RewritePatternSet &patterns,
1517  MLIRContext *context) {
1518  // TODO: remove op when offset is 0 or if it is a rotate with and
1519  // offset % shift_axis_mesh_dim_size == 0.
1520 }
1521 
1522 void ShiftOp::getAsmResultNames(
1523  function_ref<void(Value, StringRef)> setNameFn) {
1524  setNameFn(getResult(), "shift");
1525 }
1526 
1527 //===----------------------------------------------------------------------===//
1528 // mesh.update_halo op
1529 //===----------------------------------------------------------------------===//
1530 
1531 LogicalResult
1532 UpdateHaloOp::verifySymbolUses(SymbolTableCollection &symbolTable) {
1533  auto mesh = getMeshAndVerify(getOperation(), getMeshAttr(), symbolTable);
1534  if (failed(mesh)) {
1535  return failure();
1536  }
1537 
1538  return success();
1539 }
1540 
1541 //===----------------------------------------------------------------------===//
1542 // TableGen'd op method definitions
1543 //===----------------------------------------------------------------------===//
1544 
1545 #define GET_OP_CLASSES
1546 #include "mlir/Dialect/Mesh/IR/MeshOps.cpp.inc"
1547 
1548 #define GET_ATTRDEF_CLASSES
1549 #include "mlir/Dialect/Mesh/IR/MeshAttributes.cpp.inc"
1550 
1551 #define GET_TYPEDEF_CLASSES
1552 #include "mlir/Dialect/Mesh/IR/MeshTypes.cpp.inc"
1553 
1554 #include "mlir/Dialect/Mesh/IR/MeshEnums.cpp.inc"
static void copy(Location loc, Value dst, Value src, Value size, OpBuilder &builder)
Copies the given number of bytes from src to dst pointers.
static Operation * materializeConstant(Dialect *dialect, OpBuilder &builder, Attribute value, Type type, Location loc)
A utility function used to materialize a constant for a given attribute and type.
Definition: FoldUtils.cpp:50
static bool isLegalToInline(InlinerInterface &interface, Region *src, Region *insertRegion, bool shouldCloneInlinedRegion, IRMapping &valueMapping)
Utility to check that all of the operations within 'src' can be inlined.
static RankedTensorType sliceResultType(Type operandType, MeshOp mesh, ArrayRef< MeshAxis > meshAxes, int64_t sliceAxis)
Definition: MeshOps.cpp:1166
static DimensionSize operator/(DimensionSize lhs, DimensionSize rhs)
Definition: MeshOps.cpp:64
static LogicalResult verifyDimensionCompatibility(Location loc, int64_t expectedDimSize, int64_t resultDimSize, int64_t resultAxis)
Definition: MeshOps.cpp:1039
static FailureOr< MeshOp > getMeshAndVerifyAxes(Op op, SymbolTableCollection &symbolTable)
Definition: MeshOps.cpp:179
static FailureOr< MeshOp > getMeshAndVerify(Operation *op, FlatSymbolRefAttr meshSymbol, SymbolTableCollection &symbolTable)
Definition: MeshOps.cpp:127
static LogicalResult verifyScatterOrSliceOperandAndResultShape(Value operand, Value result, int64_t tensorAxis, ArrayRef< MeshAxis > meshAxes, ArrayRef< int64_t > meshShape)
Definition: MeshOps.cpp:1128
static LogicalResult verifyGatherOperandAndResultShape(Value operand, Value result, int64_t gatherAxis, ArrayRef< MeshAxis > meshAxes, ArrayRef< int64_t > meshShape)
Definition: MeshOps.cpp:1056
static LogicalResult verifyAllToAllOperandAndResultShape(Value operand, Value result, int64_t splitAxis, int64_t concatAxis, ArrayRef< MeshAxis > meshAxes, ArrayRef< int64_t > meshShape)
Definition: MeshOps.cpp:1083
static void shardShape(const InShape &inShape, const MeshShape &meshShape, const SplitAxes &splitAxes, OutShape &outShape, ArrayRef< int64_t > shardedDimsOffsets={}, ArrayRef< int64_t > haloSizes={})
Definition: MeshOps.cpp:193
static auto product(It begin, It end)
Definition: MeshOps.cpp:1028
bool isUnique(It begin, It end)
Definition: MeshOps.cpp:140
static LogicalResult verifyMeshAxes(Location loc, ArrayRef< MeshAxis > axes, MeshOp mesh)
Definition: MeshOps.cpp:156
static LogicalResult verifyInGroupDevice(Location loc, StringRef deviceName, ArrayRef< int64_t > device, Operation::operand_range deviceDynamic, ArrayRef< MeshAxis > meshAxes, ArrayRef< int64_t > meshShape)
Definition: MeshOps.cpp:1001
static Value min(ImplicitLocOpBuilder &builder, Value value, Value bound)
static ArrayRef< int64_t > getShape(Type type)
Returns the shape of the given type.
Definition: Traits.cpp:118
Attributes are known-constant values of operations.
Definition: Attributes.h:25
DenseI64ArrayAttr getDenseI64ArrayAttr(ArrayRef< int64_t > values)
Definition: Builders.cpp:163
DenseI16ArrayAttr getDenseI16ArrayAttr(ArrayRef< int16_t > values)
Definition: Builders.cpp:155
MLIRContext * getContext() const
Definition: Builders.h:56
IndexType getIndexType()
Definition: Builders.cpp:51
This is the interface that must be implemented by the dialects of operations to be inlined.
Definition: InliningUtils.h:44
DialectInlinerInterface(Dialect *dialect)
Definition: InliningUtils.h:46
A symbol reference with a reference path containing a single element.
static FlatSymbolRefAttr get(StringAttr value)
Construct a symbol reference for the given value name.
StringRef getValue() const
Returns the name of the held symbol reference.
This is a utility class for mapping one set of IR entities to another.
Definition: IRMapping.h:26
IRValueT get() const
Return the current value being used by this operand.
Definition: UseDefLists.h:160
This class coordinates rewriting a piece of IR outside of a pattern rewrite, providing a way to keep ...
Definition: PatternMatch.h:730
This class defines the main interface for locations in MLIR and acts as a non-nullable wrapper around...
Definition: Location.h:66
MLIRContext is the top-level object for a collection of MLIR operations.
Definition: MLIRContext.h:60
RAII guard to reset the insertion point of the builder when destroyed.
Definition: Builders.h:346
This class helps build Operations.
Definition: Builders.h:205
void setInsertionPoint(Block *block, Block::iterator insertPoint)
Set the insertion point to the specified location.
Definition: Builders.h:396
void setInsertionPointAfterValue(Value val)
Sets the insertion point to the node after the specified value.
Definition: Builders.h:419
Operation * create(const OperationState &state)
Creates an operation given the fields represented as an OperationState.
Definition: Builders.cpp:453
This class represents an operand of an operation.
Definition: Value.h:243
This is a value defined by a result of an operation.
Definition: Value.h:433
Location getLoc()
The source location the operation was defined or derived from.
Definition: OpDefinition.h:128
This class provides the API for a sub-set of ops that are known to be constant-like.
This provides public APIs that all operations should have.
Operation * getOperation()
Inherit getOperation from OpState.
Definition: OpDefinition.h:111
This class implements the operand iterators for the Operation class.
Definition: ValueRange.h:43
Operation is the basic unit of execution within MLIR.
Definition: Operation.h:88
bool hasTrait()
Returns true if the operation was registered with a particular trait, e.g.
Definition: Operation.h:750
void replaceUsesWithIf(ValuesT &&values, function_ref< bool(OpOperand &)> shouldReplace)
Replace uses of results of this operation with the provided values if the given callback returns true...
Definition: Operation.h:279
InFlightDiagnostic emitError(const Twine &message={})
Emit an error about fatal conditions with this operation, reporting up to any diagnostic handlers tha...
Definition: Operation.cpp:268
A special type of RewriterBase that coordinates the application of a rewrite pattern on the current I...
Definition: PatternMatch.h:749
This class contains a list of basic blocks and a link to the parent operation it is attached to.
Definition: Region.h:26
RewritePatternSet & add(ConstructorArg &&arg, ConstructorArgs &&...args)
Add an instance of each of the pattern types 'Ts' to the pattern list with the given arguments.
Definition: PatternMatch.h:811
void replaceAllUsesWith(Value from, Value to)
Find uses of from and replace them with to.
Definition: PatternMatch.h:602
virtual void eraseOp(Operation *op)
This method erases an operation that is known to have no uses.
void replaceUsesWithIf(Value from, Value to, function_ref< bool(OpOperand &)> functor, bool *allUsesReplaced=nullptr)
Find uses of from and replace them with to if the functor returns true.
void replaceAllUsesExcept(Value from, Value to, Operation *exceptedUser)
Find uses of from and replace them with to except if the user is exceptedUser.
Definition: PatternMatch.h:666
This class represents a collection of SymbolTables.
Definition: SymbolTable.h:283
Instances of the Type class are uniqued, have an immutable identifier and an optional mutable compone...
Definition: Types.h:74
This class provides an abstraction over the different types of ranges over Values.
Definition: ValueRange.h:387
This class represents an instance of an SSA value in the MLIR system, representing a computable value...
Definition: Value.h:96
Type getType() const
Return the type of this value.
Definition: Value.h:105
use_range getUses() const
Returns a range of all uses, which is useful for iterating over all uses.
Definition: Value.h:188
bool hasOneUse() const
Returns true if this value has exactly one use.
Definition: Value.h:191
Location getLoc() const
Return the location of this value.
Definition: Value.cpp:26
Operation * getDefiningOp() const
If this value is the result of an operation, return the operation that defines it.
Definition: Value.cpp:20
Base class for DenseArrayAttr that is instantiated and specialized for each supported element type be...
static DenseArrayAttrImpl get(MLIRContext *context, ArrayRef< T > content)
Builder from ArrayRef<T>.
Operation * getOwner() const
Return the owner of this operand.
Definition: UseDefLists.h:38
bool equalSplitAndPartialAxes(const MeshSharding &rhs) const
Definition: MeshOps.cpp:689
ArrayRef< int64_t > getStaticShardedDimsOffsets() const
Definition: MeshOps.h:70
bool equalHaloAndShardSizes(const MeshSharding &rhs) const
Definition: MeshOps.cpp:714
::mlir::FlatSymbolRefAttr getMeshAttr() const
Definition: MeshOps.h:64
bool equalHaloSizes(const MeshSharding &rhs) const
Definition: MeshOps.cpp:734
ArrayRef< MeshAxesAttr > getSplitAxes() const
Definition: MeshOps.h:66
bool operator!=(Value rhs) const
Definition: MeshOps.cpp:750
ReductionKind getPartialType() const
Definition: MeshOps.h:68
ArrayRef< Value > getDynamicShardedDimsOffsets() const
Definition: MeshOps.h:74
bool operator==(Value rhs) const
Definition: MeshOps.cpp:746
ArrayRef< MeshAxis > getPartialAxes() const
Definition: MeshOps.h:67
ArrayRef< Value > getDynamicHaloSizes() const
Definition: MeshOps.h:73
::llvm::StringRef getMesh() const
Definition: MeshOps.h:65
ArrayRef< int64_t > getStaticHaloSizes() const
Definition: MeshOps.h:69
MeshSharding(::mlir::FlatSymbolRefAttr mesh_=nullptr)
Definition: MeshOps.cpp:760
static MeshSharding get(::mlir::FlatSymbolRefAttr mesh_, ArrayRef< MeshAxesAttr > split_axes_, ArrayRef< MeshAxis > partial_axes_={}, ReductionKind partial_type_=ReductionKind::Sum, ArrayRef< int64_t > static_halo_sizes_={}, ArrayRef< int64_t > static_sharded_dims_offsets_={}, ArrayRef< Value > dynamic_halo_sizes_={}, ArrayRef< Value > dynamic_sharded_dims_offsets_={})
Definition: MeshOps.cpp:780
bool equalShardSizes(const MeshSharding &rhs) const
Definition: MeshOps.cpp:718
constexpr void enumerate(std::tuple< Tys... > &tuple, CallbackT &&callback)
Definition: Matchers.h:344
mesh::ReductionKind ReductionKind
int64_t collectiveProcessGroupSize(MeshAxesRange &&meshAxes, MeshShapeRange &&meshShape)
Definition: MeshOps.h:153
mesh::MeshOp getMeshOrNull(Operation *op, FlatSymbolRefAttr meshSymbol, SymbolTableCollection &symbolTableCollection)
Definition: MeshOps.h:120
void maybeInsertSourceShardingAnnotation(MeshSharding sharding, OpOperand &operand, OpBuilder &builder)
Definition: MeshOps.cpp:329
Type shardType(Type type, MeshOp mesh, MeshSharding sharding)
Definition: MeshOps.cpp:270
void maybeInsertTargetShardingAnnotation(MeshSharding sharding, OpOperand &operand, OpBuilder &builder, ShardOp &newShardOp)
Definition: MeshOps.cpp:278
ShapedType shardShapedType(ShapedType shape, MeshOp mesh, MeshSharding sharding)
Definition: MeshOps.cpp:260
int64_t shardDimension(int64_t dimSize, int64_t shardCount)
Definition: MeshOps.h:175
bool isFullReplication(MeshSharding sharding)
Definition: MeshOps.h:112
int16_t MeshAxis
Definition: MeshOps.h:26
Include the generated interface declarations.
InFlightDiagnostic emitError(Location loc)
Utility method to emit an error message using this location.
const FrozenRewritePatternSet & patterns
AffineExpr operator*(int64_t val, AffineExpr expr)
Definition: AffineExpr.h:252
void dispatchIndexOpFoldResults(ArrayRef< OpFoldResult > ofrs, SmallVectorImpl< Value > &dynamicVec, SmallVectorImpl< int64_t > &staticVec)
Helper function to dispatch multiple OpFoldResults according to the behavior of dispatchIndexOpFoldRe...
std::pair< SmallVector< int64_t >, SmallVector< Value > > decomposeMixedValues(const SmallVectorImpl< OpFoldResult > &mixedValues)
Decompose a vector of mixed static or dynamic values into the corresponding pair of arrays.
Operation * clone(OpBuilder &b, Operation *op, TypeRange newResultTypes, ValueRange newOperands)
auto get(MLIRContext *context, Ts &&...params)
Helper method that injects context only if needed, this helps unify some of the attribute constructio...
SmallVector< OpFoldResult > getMixedValues(ArrayRef< int64_t > staticValues, ValueRange dynamicValues, MLIRContext *context)
Return a vector of OpFoldResults with the same size a staticValues, but all elements for which Shaped...
LogicalResult verify(Operation *op, bool verifyRecursively=true)
Perform (potentially expensive) checks of invariants, used to detect compiler bugs,...
Definition: Verifier.cpp:423
LogicalResult foldDynamicIndexList(SmallVectorImpl< OpFoldResult > &ofrs, bool onlyNonNegative=false, bool onlyNonZero=false)
Returns "success" when any of the elements in ofrs is a constant value.
OpRewritePattern is a wrapper around RewritePattern that allows for matching and rewriting against an...
Definition: PatternMatch.h:314
This represents an operation in an abstracted form, suitable for use with the builder APIs.