Lines Matching defs:mesh

27 using namespace mlir::mesh;
105 // mesh::getMeshSharding
109 mesh::getMeshSharding(OpResult result) {
112     auto shardOp = llvm::dyn_cast<mesh::ShardOp>(user);
119     // expected to have exact one use if it has a use of `mesh.shard` without
123     auto shardOp = llvm::cast<mesh::ShardOp>(*val.getUsers().begin());
128     auto shardOp = llvm::dyn_cast<mesh::ShardOp>(user);
142       // TODO: Deduce a reasonable mesh sharding attr for def when they are
145              "only support all shard ops have the same mesh sharding attr");
153 mesh::getMeshSharding(OpOperand &opOperand) {
166 LogicalResult mesh::ShardingInterface::verifyShardingInterfaceImpl() {
208 void mesh::ShardingInterface::printLoopTypesAndIndexingMaps(raw_ostream &os) {
232                                         FlatSymbolRefAttr mesh,
235   if ((shardingOption.mesh && mesh && shardingOption.mesh != mesh) ||
248         LLVM_DEBUG(DBGS() << "sharding option conflicts because mesh axes "
254   if (mesh)
255     shardingOption.mesh = mesh;
265 mesh::detail::defaultGetShardingOption(Operation *op,
442   return MeshSharding::get(shardingOption.mesh,
478       shardingOption.mesh,
483 mesh::detail::defaultGetShardingAnnotations(
516 // To add a `mesh.shard` op for the given result, based on the details provided
530 // To add a `mesh.shard` op for the given operand, based on the details provided
547 LogicalResult mesh::detail::defaultAddShardingAnnotations(
549   assert(!shardingOption.empty && shardingOption.mesh);
559   // 1. add mesh.shard ops for all op results
567   // 2. add mesh.shard ops for all operands
605 void mesh::spmdizeFullyReplicatedOperation(
634 ShardingArray mesh::getMeshAxisAssignmentForLoopIterators(
675 bool mesh::isAtLeastOneReductionIteratorSharded(
688 SmallVector<MeshAxis> mesh::getReductionMeshAxes(
701 void mesh::spmdizeTriviallyShardableOperation(