Lines Matching defs:shardingOption

229 // Update the given `shardingOption` according to `meshAxes` and `loopIdx`
231                                         ShardingOption &shardingOption,
235   if ((shardingOption.mesh && mesh && shardingOption.mesh != mesh) ||
236       (!shardingOption.shardingArray[loopIdx].empty() &&
237        shardingOption.shardingArray[loopIdx] != meshAxes)) {
242   for (size_t i = 0; i < shardingOption.shardingArray.size(); ++i) {
247       if (llvm::is_contained(shardingOption.shardingArray[i], axis)) {
255     shardingOption.mesh = mesh;
256   if (shardingOption.shardingArray[loopIdx].empty())
257     shardingOption.shardingArray[loopIdx].append(meshAxes.begin(),
269   ShardingOption shardingOption;
277   shardingOption.shardingArray.resize(loopTypes.size());
291     // shardingOption[index]
298       if (failed(fillShardingOption(op, shardingOption, shardAttr.getMeshAttr(),
349         if (failed(fillShardingOption(op, shardingOption,
377         llvm::enumerate(shardingOption.shardingArray), [&](auto it) {
386           std::ignore = fillShardingOption(op, shardingOption, nullptr,
397   removeTrailingEmptySubArray(shardingOption.shardingArray);
399     shardingOption.empty = true;
400   return shardingOption;
404 MeshSharding getSharding(OpResult result, const ShardingOption &shardingOption,
419     if (loopIdx < shardingOption.shardingArray.size())
420       splitAxes[it.index()].append(shardingOption.shardingArray[loopIdx]);
427   for (auto it : llvm::zip(loopTypes, shardingOption.shardingArray)) {
442   return MeshSharding::get(shardingOption.mesh,
448                                            const ShardingOption &shardingOption,
463       if ((size_t)loopIdx < shardingOption.shardingArray.size() &&
464           !shardingOption.shardingArray[loopIdx].empty())
472           shardingOption.shardingArray[shardedLoopIndices[0]]);
478       shardingOption.mesh,
484     Operation *op, const ShardingOption &shardingOption) {
497         opOperand, shardingOption, maps[opOperand.getOperandNumber()]);
504     res.push_back(getSharding(result, shardingOption,
517 // in `shardingOption`, `map`, and `loopTypes`.
519                                 const ShardingOption &shardingOption,
524       getSharding(result, shardingOption, map, loopTypes, reductionLoopKinds);
531 // in `shardingOption`, `map`, and `loopTypes`.
533                                 const ShardingOption &shardingOption,
537       getSharding(opOperand, shardingOption, map);
548     Operation *op, OpBuilder &b, const ShardingOption &shardingOption) {
549   assert(!shardingOption.empty && shardingOption.mesh);
561     if (failed(addShardOp(b, result, shardingOption,
569     if (failed(addShardOp(b, opOperand, shardingOption,