| 205 | } |
| 206 | |
| 207 | Status LayoutConstraints::SetOperandLayout(const Shape& shape_with_layout, |
| 208 | const HloInstruction* instruction, |
| 209 | int64 operand_no, bool mandatory, |
| 210 | bool dfs) { |
| 211 | VLOG(3) << "SetOperandLayout : " << instruction->name() << ", operand " |
| 212 | << operand_no << " : " |
| 213 | << ShapeUtil::HumanStringWithLayout(shape_with_layout); |
| 214 | |
| 215 | const OperandLayoutConstraint* curr_shape_layout = |
| 216 | GetOperandLayoutConstraint(instruction, operand_no); |
| 217 | if (curr_shape_layout != nullptr) { |
| 218 | if (curr_shape_layout->shape_layout().MatchesLayoutInShape( |
| 219 | shape_with_layout, /*minor_to_major_only=*/true)) { |
| 220 | // New constraint matches existing constraint. Nothing to do. |
| 221 | return Status::OK(); |
| 222 | } |
| 223 | if (curr_shape_layout->mandatory()) { |
| 224 | return FailedPrecondition( |
| 225 | "Operand %d of instruction %s already has a layout constraint " |
| 226 | "%s, cannot add incompatible constraint %s", |
| 227 | operand_no, instruction->name(), |
| 228 | curr_shape_layout->shape_layout().ToString(), |
| 229 | ShapeUtil::HumanStringWithLayout(shape_with_layout)); |
| 230 | } |
| 231 | } |
| 232 | |
| 233 | // If any buffers in the operand occur in the output of the instruction, then |
| 234 | // return an error. This case is not handled because such a constraint changes |
| 235 | // layouts beyond this immediate use and is complicated to handle. |
| 236 | if (OperandBufferForwarded(instruction, operand_no)) { |
| 237 | return FailedPrecondition( |
| 238 | "Cannot constraint layout of operand %d of instruction %s " |
| 239 | "because instruction forwards operand's LogicalBuffer(s)", |
| 240 | operand_no, instruction->name()); |
| 241 | } |
| 242 | |
| 243 | auto key = std::make_pair(instruction, operand_no); |
| 244 | auto iter = operand_constraints_.find(key); |
| 245 | if (iter == operand_constraints_.end()) { |
| 246 | auto pair = std::make_pair( |
| 247 | key, OperandLayoutConstraint(ShapeLayout(shape_with_layout), |
| 248 | instruction, operand_no, mandatory, dfs)); |
| 249 | iter = operand_constraints_.insert(pair).first; |
| 250 | } else { |
| 251 | iter->second = |
| 252 | OperandLayoutConstraint(ShapeLayout(shape_with_layout), instruction, |
| 253 | operand_no, mandatory, dfs); |
| 254 | } |
| 255 | added_constraints_.push_back(&iter->second); |
| 256 | |
| 257 | return Status::OK(); |
| 258 | } |
| 259 | |
| 260 | Status LayoutConstraints::SetArrayOperandLayout( |
| 261 | const Layout& layout, const HloInstruction* instruction, int64 operand_no, |
no test coverage detected