mlir/lib/Dialect/Mesh/Transforms/ShardingPropagation.cpp - llvm-project - Git at Google

 //===- ShardingPropagation.cpp ------------------------------------- C++ --===//
 //
 // 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/Transforms/Passes.h"

 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/Dialect/Mesh/IR/MeshDialect.h"
 #include "mlir/Dialect/Mesh/IR/MeshOps.h"
 #include "mlir/Dialect/Mesh/Interfaces/ShardingInterface.h"
 #include "mlir/IR/Verifier.h"
 #include "mlir/Interfaces/FunctionInterfaces.h"
 #include "mlir/Pass/Pass.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/iterator_range.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/raw_ostream.h"
 #include <algorithm>
 #include <vector>

 namespace mlir {
 namespace mesh {
 #define GEN_PASS_DEF_SHARDINGPROPAGATION
 #include "mlir/Dialect/Mesh/Transforms/Passes.h.inc"
 } // namespace mesh
 } // namespace mlir

 #define DEBUG_TYPE "sharding-propagation"
 #define DBGS() (llvm::dbgs() << "[" DEBUG_TYPE << "]: ")

 using namespace mlir;
 using namespace mlir::mesh;

 enum class ReshardingRquirementKind {
   NO_RESHARDING = 0,
   NO_RESHARDING_FOR_EXPLICIT_ANNOTATIONS,
   RESHARDING_FOR_EXPLICIT_ANNOTATIONS
 };

 #ifdef LLVM_DEBUG

 template <typename T>
 static llvm::raw_ostream &operator<<(llvm::raw_ostream &stream,
                                      const SmallVector<T> &vec);
 template <typename... Ts>
 static llvm::raw_ostream &operator<<(llvm::raw_ostream &stream,
                                      const std::tuple<Ts...> &t);
 static llvm::raw_ostream &operator<<(llvm::raw_ostream &stream,
                                      ReshardingRquirementKind v);

 template <typename Stream, typename Range>
 static Stream &printRange(Stream &stream, Range &&range) {
   stream << "[";
   llvm::for_each(range, [&stream](auto &v) {
     stream << v;
     stream << ", ";
   });
   return stream << "]";
 }

 template <typename T>
 static llvm::raw_ostream &operator<<(llvm::raw_ostream &stream,
                                      const SmallVector<T> &vec) {
   return printRange(stream, vec);
 }

 [[maybe_unused]] static llvm::raw_ostream &operator<<(llvm::raw_ostream &stream,
                                                       const ShardingOption &v) {
   return stream << "{empty = " << v.empty << ", mesh" << v.mesh
                 << ", shardingArray = " << v.shardingArray << "}";
 }

 template <typename Stream, typename... Ts, size_t... Is>
 static Stream &printTuple(Stream &stream, std::tuple<Ts...> tuple,
                           std::index_sequence<Is...>) {
   static_assert(sizeof...(Is) == sizeof...(Ts),
                 "Indices must have same number of elements as tuple types!");
   static_assert(sizeof...(Ts) > 0, "Cannot insert empty tuple into stream.");

   stream << "{";
   ((stream << std::get<Is>(tuple) << ", "), ...);
   return stream << "}";
 }

 template <typename... Ts>
 static llvm::raw_ostream &operator<<(llvm::raw_ostream &stream,
                                      const std::tuple<Ts...> &t) {
   return printTuple(stream, t, std::index_sequence_for<Ts...>{});
 }

 [[maybe_unused]] static llvm::raw_ostream &
 operator<<(llvm::raw_ostream &stream, ReshardingRquirementKind v) {
   return stream << static_cast<int>(v);
 }

 #endif // LLVM_DEBUG

 //===----------------------------------------------------------------------===//
 // Utilities
 //===----------------------------------------------------------------------===//

 // This method retrieves all potential sharding attributes, prioritizing
 // specific shardings. For example, mustShardings = [shard0, None] and
 // optionalShardings = [None, shard1], the result will be [[shard0, shard1],
 // [shard0, None]]
 static SmallVector<std::vector<MeshSharding>>
 getOrderedPossibleShardingAttrs(ArrayRef<MeshSharding> mustShardings,
                                 ArrayRef<MeshSharding> optionalShardings) {
   SmallVector<std::vector<MeshSharding>> allShardingAttrs;
   std::vector<MeshSharding> curShardingAttrs;

   std::function<void(size_t)> dfsCreateShardingAttrs = [&](size_t i) {
     if (i == mustShardings.size()) {
       allShardingAttrs.push_back(std::vector<MeshSharding>(curShardingAttrs));
       return;
     }

     if (mustShardings[i]) {
       curShardingAttrs.push_back(mustShardings[i]);
       dfsCreateShardingAttrs(i + 1);
       curShardingAttrs.pop_back();
       return;
     }

     if (optionalShardings[i]) {
       curShardingAttrs.push_back(optionalShardings[i]);
       dfsCreateShardingAttrs(i + 1);
       curShardingAttrs.pop_back();
       curShardingAttrs.push_back({});
       dfsCreateShardingAttrs(i + 1);
       curShardingAttrs.pop_back();
       return;
     }

     curShardingAttrs.push_back({});
     dfsCreateShardingAttrs(i + 1);
     curShardingAttrs.pop_back();
   };

   dfsCreateShardingAttrs(0);
   return allShardingAttrs;
 }

 // The order of preference is form highest to lowest:
 // 1. No resharding is required (all existing annotations are compatible).
 // 2. No resharding for operands/results that have annotation specifically
 //   targeting this operation. This means
 //   * operands that are the result of `mesh.shard` ops marked with
 //     `annotate_for_users`.
 //   * results that are annotated with `mesh.shard` ops without
 //     `annotate_for_users`.
 // 3. All other cases. Resharding is required for operands/results with
 //   annotation targeting explicitly this operation.
 ReshardingRquirementKind getReshardingRquirementKind(
     Operation *op, const std::vector<MeshSharding> &operandAndResultShardings) {
   ReshardingRquirementKind res = ReshardingRquirementKind::NO_RESHARDING;

   size_t operandsCount = op->getOperands().size();
   auto operandShardings =
       llvm::make_range(operandAndResultShardings.begin(),
                        operandAndResultShardings.begin() + operandsCount);
   auto resultShardings =
       llvm::make_range(operandAndResultShardings.begin() + operandsCount,
                        operandAndResultShardings.end());

   for (auto [operand, sharding] :
        llvm::zip_equal(op->getOperands(), operandShardings)) {
     ShardOp shardOp = llvm::dyn_cast_or_null<ShardOp>(operand.getDefiningOp());
     if (!shardOp) {
       continue;
     }
     bool needsResharding = sharding != shardOp.getSharding();
     bool isExplicitAnnotationForThisOp = shardOp.getAnnotateForUsers();
     if (needsResharding) {
       if (isExplicitAnnotationForThisOp) {
         // This is the worst case. No need to continue.
         return ReshardingRquirementKind::RESHARDING_FOR_EXPLICIT_ANNOTATIONS;
       }
       res = ReshardingRquirementKind::NO_RESHARDING_FOR_EXPLICIT_ANNOTATIONS;
     }
   }

   for (auto [result, sharding] :
        llvm::zip_equal(op->getResults(), resultShardings)) {
     for (auto user : result.getUsers()) {
       ShardOp shardOp = llvm::dyn_cast<ShardOp>(user);
       if (!shardOp) {
         continue;
       }
       bool needsResharding = sharding != shardOp.getSharding();
       bool isExplicitAnnotationForThisOp = !shardOp.getAnnotateForUsers();
       if (needsResharding) {
         if (isExplicitAnnotationForThisOp) {
           // This is the worst case. No need to continue.
           return ReshardingRquirementKind::RESHARDING_FOR_EXPLICIT_ANNOTATIONS;
         }
         res = ReshardingRquirementKind::NO_RESHARDING_FOR_EXPLICIT_ANNOTATIONS;
       }
     }
   }

   return res;
 }

 // From all the operand and result sharding combinations,
 // return the one that is most desirable.
 // The order of preference is:
 // 1. No resharding with respect to existing sharding annotations.
 // 2. Resharding for values that have already annotations that do not target
 //    this op.
 // 3. Resharding of existing explicit sharding annotations for this op.
 static FailureOr<ShardingOption> selectShardingOption(
     ShardingInterface shardingOp,
     ArrayRef<std::vector<MeshSharding>> possibleOperandShardingAttrs,
     ArrayRef<std::vector<MeshSharding>> possibleResultShardingAttrs) {
   SmallVector<std::tuple<ShardingOption, ReshardingRquirementKind>>
       shardingOptionsAndReshardingRequirements;

   for (ArrayRef<MeshSharding> resultShardings : possibleResultShardingAttrs) {
     for (ArrayRef<MeshSharding> operandShardings :
          possibleOperandShardingAttrs) {
       FailureOr<ShardingOption> shardingOption =
           shardingOp.getShardingOption(operandShardings, resultShardings);
       if (failed(shardingOption) || shardingOption->empty) {
         continue;
       }
       // These shardings may not be the same as those in operandShardings and
       // resultShardings.
       // They may be missing some annotations.
       // Whatever is returned by getShardingAnnotations is exactly what the op
       // needs.
       FailureOr<std::vector<MeshSharding>> operandAndResultShardings =
           shardingOp.getShardingAnnotations(*shardingOption);
       if (failed(operandAndResultShardings)) {
         return failure();
       }

       // LLVM_DEBUG(DBGS() << "operandAndResultShardings = "
       //                   << *operandAndResultShardings << "\n";);

       ReshardingRquirementKind reshardingRquirement =
           getReshardingRquirementKind(shardingOp, *operandAndResultShardings);
       if (reshardingRquirement == ReshardingRquirementKind::NO_RESHARDING) {
         // This is the best case. No need to go on.
         return *shardingOption;
       }

       shardingOptionsAndReshardingRequirements.emplace_back(
           std::move(*shardingOption), reshardingRquirement);
     }
   }

   if (shardingOptionsAndReshardingRequirements.empty()) {
     return ShardingOption::makeEmpty();
   }

   std::partial_sort(
       shardingOptionsAndReshardingRequirements.begin(),
       shardingOptionsAndReshardingRequirements.begin() + 1,
       shardingOptionsAndReshardingRequirements.end(),
       [](const std::tuple<ShardingOption, ReshardingRquirementKind> &a,
          const std::tuple<ShardingOption, ReshardingRquirementKind> &b) {
         return std::get<ReshardingRquirementKind>(a) <
                std::get<ReshardingRquirementKind>(b);
       });

   LLVM_DEBUG(DBGS() << "shardingOptionsAndReshardingRequirements = "
                     << shardingOptionsAndReshardingRequirements << "\n";);

   return std::get<ShardingOption>(
       shardingOptionsAndReshardingRequirements.front());
 }

 // For each operation that implements the ShardingInterface, infer the sharding
 // option of the operation from its operands and/or results using the
 // `getShardingOption` method. If the inferred sharding option is not empty, add
 // a `mesh.shard` operation for all remaining operands and results that do not
 // have sharding annotations.
 static LogicalResult visitOp(Operation *op, OpBuilder &builder) {
   if (op->hasTrait<OpTrait::IsTerminator>() ||
       llvm::isa<mesh::ShardOp, mesh::ShardingOp>(op))
     return success();

   ShardingInterface shardingOp = llvm::dyn_cast<ShardingInterface>(op);
   if (!shardingOp) {
     op->emitOpError() << "sharding interface is not implemented.";
     return failure();
   }

   // collect MeshSharding from results
   std::vector<MeshSharding> allowConflictsResultShardings;
   allowConflictsResultShardings.resize(op->getNumResults());
   std::vector<MeshSharding> resultMustShardings;
   resultMustShardings.resize(op->getNumResults());
   for (OpResult result : op->getResults()) {
     FailureOr<std::pair<bool, MeshSharding>> maybeShardAttr =
         getMeshSharding(result);
     if (failed(maybeShardAttr))
       continue;
     if (!maybeShardAttr->first)
       resultMustShardings[result.getResultNumber()] = maybeShardAttr->second;
     else
       allowConflictsResultShardings[result.getResultNumber()] =
           maybeShardAttr->second;
   }

   // collect MeshSharding from operands
   std::vector<MeshSharding> allowConflictsOperandShardings;
   allowConflictsOperandShardings.resize(op->getNumOperands());
   std::vector<MeshSharding> operandMustShardings;
   operandMustShardings.resize(op->getNumOperands());
   for (OpOperand &opOperand : op->getOpOperands()) {
     FailureOr<std::pair<bool, MeshSharding>> maybeShardAttr =
         getMeshSharding(opOperand);
     if (failed(maybeShardAttr))
       continue;

     if (maybeShardAttr->first)
       operandMustShardings[opOperand.getOperandNumber()] =
           maybeShardAttr->second;
     else
       allowConflictsOperandShardings[opOperand.getOperandNumber()] =
           maybeShardAttr->second;
   }

   // try to get the sharding option
   SmallVector<std::vector<MeshSharding>> possibleOperandShardingAttrs =
       getOrderedPossibleShardingAttrs(operandMustShardings,
                                       allowConflictsOperandShardings);
   SmallVector<std::vector<MeshSharding>> possibleResultShardingAttrs =
       getOrderedPossibleShardingAttrs(resultMustShardings,
                                       allowConflictsResultShardings);
   FailureOr<ShardingOption> shardingOption = selectShardingOption(
       shardingOp, possibleOperandShardingAttrs, possibleResultShardingAttrs);

   if (failed(shardingOption)) {
     op->emitOpError() << "fail to get sharding option.";
     return failure();
   }

   LLVM_DEBUG(DBGS() << "Selected sharding option: " << *shardingOption << "\n");

   // sharding info is empty, return immediately
   if (shardingOption->empty)
     return success();

   if (failed(shardingOp.addShardingAnnotations(builder, *shardingOption))) {
     op->emitOpError() << "fail to set sharding annotations.";
     return failure();
   }
   return success();
 }

 //===----------------------------------------------------------------------===//
 // ShardingPropagation
 //===----------------------------------------------------------------------===//
 struct ShardingPropagation
     : public mesh::impl::ShardingPropagationBase<ShardingPropagation> {
   void runOnOperation() override {
     FunctionOpInterface funcOp = getOperation();
     MLIRContext *ctx = funcOp.getContext();
     Region &region = funcOp.getFunctionBody();
     OpBuilder builder(ctx);
     if (!region.hasOneBlock()) {
       funcOp.emitOpError() << "only one block is supported!";
       signalPassFailure();
     }
     Block &block = region.front();

     LLVM_DEBUG(
         DBGS() << "print all the ops' iterator types and indexing maps in the "
                   "block.\n";
         for (Operation &op
              : block.getOperations()) {
           if (auto shardingOp = llvm::dyn_cast<ShardingInterface>(&op))
             shardingOp.printLoopTypesAndIndexingMaps(llvm::dbgs());
         });

     // 1. propagate in reversed order
     for (Operation &op : llvm::make_early_inc_range(llvm::reverse(block)))
       if (failed(visitOp(&op, builder)))
         return signalPassFailure();

     LLVM_DEBUG(DBGS() << "After reversed order propagation:\n"
                       << funcOp << "\n");
     LLVM_DEBUG(assert(succeeded(mlir::verify(funcOp))));

     // 2. propagate in original order
     for (Operation &op : llvm::make_early_inc_range(block))
       if (failed(visitOp(&op, builder)))
         return signalPassFailure();
   }
 };
	//===- ShardingPropagation.cpp ------------------------------------- C++ --===//
	//
	// 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/Transforms/Passes.h"

	#include "mlir/Dialect/Func/IR/FuncOps.h"
	#include "mlir/Dialect/Mesh/IR/MeshDialect.h"
	#include "mlir/Dialect/Mesh/IR/MeshOps.h"
	#include "mlir/Dialect/Mesh/Interfaces/ShardingInterface.h"
	#include "mlir/IR/Verifier.h"
	#include "mlir/Interfaces/FunctionInterfaces.h"
	#include "mlir/Pass/Pass.h"
	#include "llvm/ADT/STLExtras.h"
	#include "llvm/ADT/SmallVector.h"
	#include "llvm/ADT/iterator_range.h"
	#include "llvm/Support/Debug.h"
	#include "llvm/Support/raw_ostream.h"
	#include <algorithm>
	#include <vector>

	namespace mlir {
	namespace mesh {
	#define GEN_PASS_DEF_SHARDINGPROPAGATION
	#include "mlir/Dialect/Mesh/Transforms/Passes.h.inc"
	} // namespace mesh
	} // namespace mlir

	#define DEBUG_TYPE "sharding-propagation"
	#define DBGS() (llvm::dbgs() << "[" DEBUG_TYPE << "]: ")

	using namespace mlir;
	using namespace mlir::mesh;

	enum class ReshardingRquirementKind {
	NO_RESHARDING = 0,
	NO_RESHARDING_FOR_EXPLICIT_ANNOTATIONS,
	RESHARDING_FOR_EXPLICIT_ANNOTATIONS
	};

	#ifdef LLVM_DEBUG

	template <typename T>
	static llvm::raw_ostream &operator<<(llvm::raw_ostream &stream,
	const SmallVector<T> &vec);
	template <typename... Ts>
	static llvm::raw_ostream &operator<<(llvm::raw_ostream &stream,
	const std::tuple<Ts...> &t);
	static llvm::raw_ostream &operator<<(llvm::raw_ostream &stream,
	ReshardingRquirementKind v);

	template <typename Stream, typename Range>
	static Stream &printRange(Stream &stream, Range &&range) {
	stream << "[";
	llvm::for_each(range, [&stream](auto &v) {
	stream << v;
	stream << ", ";
	});
	return stream << "]";
	}

	template <typename T>
	static llvm::raw_ostream &operator<<(llvm::raw_ostream &stream,
	const SmallVector<T> &vec) {
	return printRange(stream, vec);
	}

	[[maybe_unused]] static llvm::raw_ostream &operator<<(llvm::raw_ostream &stream,
	const ShardingOption &v) {
	return stream << "{empty = " << v.empty << ", mesh" << v.mesh
	<< ", shardingArray = " << v.shardingArray << "}";
	}

	template <typename Stream, typename... Ts, size_t... Is>
	static Stream &printTuple(Stream &stream, std::tuple<Ts...> tuple,
	std::index_sequence<Is...>) {
	static_assert(sizeof...(Is) == sizeof...(Ts),
	"Indices must have same number of elements as tuple types!");
	static_assert(sizeof...(Ts) > 0, "Cannot insert empty tuple into stream.");

	stream << "{";
	((stream << std::get<Is>(tuple) << ", "), ...);
	return stream << "}";
	}

	template <typename... Ts>
	static llvm::raw_ostream &operator<<(llvm::raw_ostream &stream,
	const std::tuple<Ts...> &t) {
	return printTuple(stream, t, std::index_sequence_for<Ts...>{});
	}

	[[maybe_unused]] static llvm::raw_ostream &
	operator<<(llvm::raw_ostream &stream, ReshardingRquirementKind v) {
	return stream << static_cast<int>(v);
	}

	#endif // LLVM_DEBUG

	//===----------------------------------------------------------------------===//
	// Utilities
	//===----------------------------------------------------------------------===//

	// This method retrieves all potential sharding attributes, prioritizing
	// specific shardings. For example, mustShardings = [shard0, None] and
	// optionalShardings = [None, shard1], the result will be [[shard0, shard1],
	// [shard0, None]]
	static SmallVector<std::vector<MeshSharding>>
	getOrderedPossibleShardingAttrs(ArrayRef<MeshSharding> mustShardings,
	ArrayRef<MeshSharding> optionalShardings) {
	SmallVector<std::vector<MeshSharding>> allShardingAttrs;
	std::vector<MeshSharding> curShardingAttrs;

	std::function<void(size_t)> dfsCreateShardingAttrs = [&](size_t i) {
	if (i == mustShardings.size()) {
	allShardingAttrs.push_back(std::vector<MeshSharding>(curShardingAttrs));
	return;
	}

	if (mustShardings[i]) {
	curShardingAttrs.push_back(mustShardings[i]);
	dfsCreateShardingAttrs(i + 1);
	curShardingAttrs.pop_back();
	return;
	}

	if (optionalShardings[i]) {
	curShardingAttrs.push_back(optionalShardings[i]);
	dfsCreateShardingAttrs(i + 1);
	curShardingAttrs.pop_back();
	curShardingAttrs.push_back({});
	dfsCreateShardingAttrs(i + 1);
	curShardingAttrs.pop_back();
	return;
	}

	curShardingAttrs.push_back({});
	dfsCreateShardingAttrs(i + 1);
	curShardingAttrs.pop_back();
	};

	dfsCreateShardingAttrs(0);
	return allShardingAttrs;
	}

	// The order of preference is form highest to lowest:
	// 1. No resharding is required (all existing annotations are compatible).
	// 2. No resharding for operands/results that have annotation specifically
	// targeting this operation. This means
	// * operands that are the result of `mesh.shard` ops marked with
	// `annotate_for_users`.
	// * results that are annotated with `mesh.shard` ops without
	// `annotate_for_users`.
	// 3. All other cases. Resharding is required for operands/results with
	// annotation targeting explicitly this operation.
	ReshardingRquirementKind getReshardingRquirementKind(
	Operation *op, const std::vector<MeshSharding> &operandAndResultShardings) {
	ReshardingRquirementKind res = ReshardingRquirementKind::NO_RESHARDING;

	size_t operandsCount = op->getOperands().size();
	auto operandShardings =
	llvm::make_range(operandAndResultShardings.begin(),
	operandAndResultShardings.begin() + operandsCount);
	auto resultShardings =
	llvm::make_range(operandAndResultShardings.begin() + operandsCount,
	operandAndResultShardings.end());

	for (auto [operand, sharding] :
	llvm::zip_equal(op->getOperands(), operandShardings)) {
	ShardOp shardOp = llvm::dyn_cast_or_null<ShardOp>(operand.getDefiningOp());
	if (!shardOp) {
	continue;
	}
	bool needsResharding = sharding != shardOp.getSharding();
	bool isExplicitAnnotationForThisOp = shardOp.getAnnotateForUsers();
	if (needsResharding) {
	if (isExplicitAnnotationForThisOp) {
	// This is the worst case. No need to continue.
	return ReshardingRquirementKind::RESHARDING_FOR_EXPLICIT_ANNOTATIONS;
	}
	res = ReshardingRquirementKind::NO_RESHARDING_FOR_EXPLICIT_ANNOTATIONS;
	}
	}

	for (auto [result, sharding] :
	llvm::zip_equal(op->getResults(), resultShardings)) {
	for (auto user : result.getUsers()) {
	ShardOp shardOp = llvm::dyn_cast<ShardOp>(user);
	if (!shardOp) {
	continue;
	}
	bool needsResharding = sharding != shardOp.getSharding();
	bool isExplicitAnnotationForThisOp = !shardOp.getAnnotateForUsers();
	if (needsResharding) {
	if (isExplicitAnnotationForThisOp) {
	// This is the worst case. No need to continue.
	return ReshardingRquirementKind::RESHARDING_FOR_EXPLICIT_ANNOTATIONS;
	}
	res = ReshardingRquirementKind::NO_RESHARDING_FOR_EXPLICIT_ANNOTATIONS;
	}
	}
	}

	return res;
	}

	// From all the operand and result sharding combinations,
	// return the one that is most desirable.
	// The order of preference is:
	// 1. No resharding with respect to existing sharding annotations.
	// 2. Resharding for values that have already annotations that do not target
	// this op.
	// 3. Resharding of existing explicit sharding annotations for this op.
	static FailureOr<ShardingOption> selectShardingOption(
	ShardingInterface shardingOp,
	ArrayRef<std::vector<MeshSharding>> possibleOperandShardingAttrs,
	ArrayRef<std::vector<MeshSharding>> possibleResultShardingAttrs) {
	SmallVector<std::tuple<ShardingOption, ReshardingRquirementKind>>
	shardingOptionsAndReshardingRequirements;

	for (ArrayRef<MeshSharding> resultShardings : possibleResultShardingAttrs) {
	for (ArrayRef<MeshSharding> operandShardings :
	possibleOperandShardingAttrs) {
	FailureOr<ShardingOption> shardingOption =
	shardingOp.getShardingOption(operandShardings, resultShardings);
	if (failed(shardingOption) \|\| shardingOption->empty) {
	continue;
	}
	// These shardings may not be the same as those in operandShardings and
	// resultShardings.
	// They may be missing some annotations.
	// Whatever is returned by getShardingAnnotations is exactly what the op
	// needs.
	FailureOr<std::vector<MeshSharding>> operandAndResultShardings =
	shardingOp.getShardingAnnotations(*shardingOption);
	if (failed(operandAndResultShardings)) {
	return failure();
	}

	// LLVM_DEBUG(DBGS() << "operandAndResultShardings = "
	// << *operandAndResultShardings << "\n";);

	ReshardingRquirementKind reshardingRquirement =
	getReshardingRquirementKind(shardingOp, *operandAndResultShardings);
	if (reshardingRquirement == ReshardingRquirementKind::NO_RESHARDING) {
	// This is the best case. No need to go on.
	return *shardingOption;
	}

	shardingOptionsAndReshardingRequirements.emplace_back(
	std::move(*shardingOption), reshardingRquirement);
	}
	}

	if (shardingOptionsAndReshardingRequirements.empty()) {
	return ShardingOption::makeEmpty();
	}

	std::partial_sort(
	shardingOptionsAndReshardingRequirements.begin(),
	shardingOptionsAndReshardingRequirements.begin() + 1,
	shardingOptionsAndReshardingRequirements.end(),
	[](const std::tuple<ShardingOption, ReshardingRquirementKind> &a,
	const std::tuple<ShardingOption, ReshardingRquirementKind> &b) {
	return std::get<ReshardingRquirementKind>(a) <
	std::get<ReshardingRquirementKind>(b);
	});

	LLVM_DEBUG(DBGS() << "shardingOptionsAndReshardingRequirements = "
	<< shardingOptionsAndReshardingRequirements << "\n";);

	return std::get<ShardingOption>(
	shardingOptionsAndReshardingRequirements.front());
	}

	// For each operation that implements the ShardingInterface, infer the sharding
	// option of the operation from its operands and/or results using the
	// `getShardingOption` method. If the inferred sharding option is not empty, add
	// a `mesh.shard` operation for all remaining operands and results that do not
	// have sharding annotations.
	static LogicalResult visitOp(Operation *op, OpBuilder &builder) {
	if (op->hasTrait<OpTrait::IsTerminator>() \|\|
	llvm::isa<mesh::ShardOp, mesh::ShardingOp>(op))
	return success();

	ShardingInterface shardingOp = llvm::dyn_cast<ShardingInterface>(op);
	if (!shardingOp) {
	op->emitOpError() << "sharding interface is not implemented.";
	return failure();
	}

	// collect MeshSharding from results
	std::vector<MeshSharding> allowConflictsResultShardings;
	allowConflictsResultShardings.resize(op->getNumResults());
	std::vector<MeshSharding> resultMustShardings;
	resultMustShardings.resize(op->getNumResults());
	for (OpResult result : op->getResults()) {
	FailureOr<std::pair<bool, MeshSharding>> maybeShardAttr =
	getMeshSharding(result);
	if (failed(maybeShardAttr))
	continue;
	if (!maybeShardAttr->first)
	resultMustShardings[result.getResultNumber()] = maybeShardAttr->second;
	else
	allowConflictsResultShardings[result.getResultNumber()] =
	maybeShardAttr->second;
	}

	// collect MeshSharding from operands
	std::vector<MeshSharding> allowConflictsOperandShardings;
	allowConflictsOperandShardings.resize(op->getNumOperands());
	std::vector<MeshSharding> operandMustShardings;
	operandMustShardings.resize(op->getNumOperands());
	for (OpOperand &opOperand : op->getOpOperands()) {
	FailureOr<std::pair<bool, MeshSharding>> maybeShardAttr =
	getMeshSharding(opOperand);
	if (failed(maybeShardAttr))
	continue;

	if (maybeShardAttr->first)
	operandMustShardings[opOperand.getOperandNumber()] =
	maybeShardAttr->second;
	else
	allowConflictsOperandShardings[opOperand.getOperandNumber()] =
	maybeShardAttr->second;
	}

	// try to get the sharding option
	SmallVector<std::vector<MeshSharding>> possibleOperandShardingAttrs =
	getOrderedPossibleShardingAttrs(operandMustShardings,
	allowConflictsOperandShardings);
	SmallVector<std::vector<MeshSharding>> possibleResultShardingAttrs =
	getOrderedPossibleShardingAttrs(resultMustShardings,
	allowConflictsResultShardings);
	FailureOr<ShardingOption> shardingOption = selectShardingOption(
	shardingOp, possibleOperandShardingAttrs, possibleResultShardingAttrs);

	if (failed(shardingOption)) {
	op->emitOpError() << "fail to get sharding option.";
	return failure();
	}

	LLVM_DEBUG(DBGS() << "Selected sharding option: " << *shardingOption << "\n");

	// sharding info is empty, return immediately
	if (shardingOption->empty)
	return success();

	if (failed(shardingOp.addShardingAnnotations(builder, *shardingOption))) {
	op->emitOpError() << "fail to set sharding annotations.";
	return failure();
	}
	return success();
	}

	//===----------------------------------------------------------------------===//
	// ShardingPropagation
	//===----------------------------------------------------------------------===//
	struct ShardingPropagation
	: public mesh::impl::ShardingPropagationBase<ShardingPropagation> {
	void runOnOperation() override {
	FunctionOpInterface funcOp = getOperation();
	MLIRContext *ctx = funcOp.getContext();
	Region &region = funcOp.getFunctionBody();
	OpBuilder builder(ctx);
	if (!region.hasOneBlock()) {
	funcOp.emitOpError() << "only one block is supported!";
	signalPassFailure();
	}
	Block &block = region.front();

	LLVM_DEBUG(
	DBGS() << "print all the ops' iterator types and indexing maps in the "
	"block.\n";
	for (Operation &op
	: block.getOperations()) {
	if (auto shardingOp = llvm::dyn_cast<ShardingInterface>(&op))
	shardingOp.printLoopTypesAndIndexingMaps(llvm::dbgs());
	});

	// 1. propagate in reversed order
	for (Operation &op : llvm::make_early_inc_range(llvm::reverse(block)))
	if (failed(visitOp(&op, builder)))
	return signalPassFailure();

	LLVM_DEBUG(DBGS() << "After reversed order propagation:\n"
	<< funcOp << "\n");
	LLVM_DEBUG(assert(succeeded(mlir::verify(funcOp))));

	// 2. propagate in original order
	for (Operation &op : llvm::make_early_inc_range(block))
	if (failed(visitOp(&op, builder)))
	return signalPassFailure();
	}
	};