| 156 | } |
| 157 | |
| 158 | Status LayoutConstraints::SetBufferLayout(const Layout& layout, |
| 159 | const LogicalBuffer& buffer, |
| 160 | bool mandatory, bool dfs) { |
| 161 | VLOG(3) << "SetBufferLayout : " << buffer << " : " |
| 162 | << LayoutUtil::HumanString(layout); |
| 163 | |
| 164 | TF_RETURN_IF_ERROR(points_to_analysis_.VerifyBuffer(buffer)); |
| 165 | if (!buffer.IsArray()) { |
| 166 | return FailedPrecondition( |
| 167 | "Layout of buffer %s cannot be constrained because buffer is not " |
| 168 | "array-shaped, has shape: %s", |
| 169 | buffer.ToString(), ShapeUtil::HumanString(buffer.shape())); |
| 170 | } |
| 171 | TF_RETURN_IF_ERROR( |
| 172 | LayoutUtil::ValidateLayoutForShape(layout, buffer.shape())); |
| 173 | |
| 174 | auto iter = buffer_constraints_.find(&buffer); |
| 175 | if (iter != buffer_constraints_.end()) { |
| 176 | const BufferLayoutConstraint& curr_constraint = iter->second; |
| 177 | if (Layout::Equal().MinorToMajorOnly()(curr_constraint.layout(), layout)) { |
| 178 | // New constraint matches existing constraint. Nothing to do. |
| 179 | return Status::OK(); |
| 180 | } |
| 181 | if (curr_constraint.mandatory()) { |
| 182 | if (!mandatory) { |
| 183 | VLOG(3) << "Buffer" << buffer |
| 184 | << " already has a mandatory layout constrain, skipping"; |
| 185 | return Status::OK(); |
| 186 | } |
| 187 | return FailedPrecondition( |
| 188 | "Buffer %s already has the layout constraint %s, cannot add " |
| 189 | "incompatible constraint %s", |
| 190 | buffer.ToString(), LayoutUtil::HumanString(curr_constraint.layout()), |
| 191 | LayoutUtil::HumanString(layout)); |
| 192 | } |
| 193 | iter->second = BufferLayoutConstraint(layout, buffer, mandatory, dfs); |
| 194 | } else { |
| 195 | TF_RET_CHECK(unconstrained_buffer_ids_.erase(buffer.id()) == 1) |
| 196 | << buffer.ToString(); |
| 197 | iter = buffer_constraints_ |
| 198 | .insert(std::make_pair( |
| 199 | &buffer, |
| 200 | BufferLayoutConstraint(layout, buffer, mandatory, dfs))) |
| 201 | .first; |
| 202 | } |
| 203 | added_constraints_.push_back(&iter->second); |
| 204 | return Status::OK(); |
| 205 | } |
| 206 | |
| 207 | Status LayoutConstraints::SetOperandLayout(const Shape& shape_with_layout, |
| 208 | const HloInstruction* instruction, |
no test coverage detected