Lines Matching defs:shardOp

630         ShardOp shardOp = llvm::dyn_cast<ShardOp>(useOp);
631         assert(shardOp);
632         MeshOp mesh = getMesh(shardOp, symbolTableCollection);
634                                           shardOp.getSharding()));
688     ShardOp shardOp = llvm::cast<ShardOp>(definingOp);
689     return MeshSharding(shardOp.getSharding());
709                     ShardOp shardOp = llvm::cast<ShardOp>(userOp);
710                     return MeshSharding(shardOp.getSharding());
716 spmdizeOperation(ShardOp shardOp, IRMapping &spmdizationMap,
724       dyn_cast_or_null<ShardOp>(shardOp.getSrc().getDefiningOp());
726     targetSpmdValue = spmdizationMap.lookup(shardOp.getSrc());
731     targetSpmdValue = reshard(builder, srcShardOp, shardOp, srcSpmdValue,
735   assert(!spmdizationMap.contains(shardOp.getResult()));
736   spmdizationMap.map(shardOp.getResult(), targetSpmdValue);
748   ShardOp shardOp = llvm::dyn_cast<ShardOp>(op);
749   if (shardOp) {
750     return spmdizeOperation(shardOp, spmdizationMap, symbolTableCollection,