| 2214 | } |
| 2215 | |
| 2216 | Status AlgebraicSimplifierVisitor::HandleMultiply(HloInstruction* multiply) { |
| 2217 | HloInstruction *lhs, *rhs; |
| 2218 | CHECK(Match(multiply, m::Multiply(m::Op(&lhs), m::Op(&rhs)))); |
| 2219 | // LHS*1 => LHS |
| 2220 | VLOG(10) << "trying transform [LHS*1 => LHS]: " << multiply->ToString(); |
| 2221 | if (IsAll(rhs, 1) && ReplaceInstructionIfSameShape(multiply, lhs)) { |
| 2222 | return Status::OK(); |
| 2223 | } |
| 2224 | // 1*RHS => RHS |
| 2225 | VLOG(10) << "trying transform [1*RHS => RHS]: " << multiply->ToString(); |
| 2226 | if (IsAll(lhs, 1) && ReplaceInstructionIfSameShape(multiply, rhs)) { |
| 2227 | return Status::OK(); |
| 2228 | } |
| 2229 | |
| 2230 | // 0*RHS => 0. Only applies for integral types for correct NaN-handling. |
| 2231 | if (IsAll(lhs, 0) && |
| 2232 | primitive_util::IsIntegralType(multiply->shape().element_type()) && |
| 2233 | ReplaceInstructionIfSameShape(multiply, lhs)) { |
| 2234 | return Status::OK(); |
| 2235 | } |
| 2236 | // LHS*0 => 0 |
| 2237 | if (IsAll(rhs, 0) && |
| 2238 | primitive_util::IsIntegralType(multiply->shape().element_type()) && |
| 2239 | ReplaceInstructionIfSameShape(multiply, rhs)) { |
| 2240 | return Status::OK(); |
| 2241 | } |
| 2242 | |
| 2243 | VLOG(10) << "trying transform [(A * C1) * C2 => A * (C1 * C2)]"; |
| 2244 | HloInstruction *a, *c1, *c2; |
| 2245 | if (Match(multiply, |
| 2246 | m::Multiply(m::Multiply(m::NonConstant(&a), m::Constant(&c1)), |
| 2247 | m::Constant(&c2))) || |
| 2248 | Match(multiply, |
| 2249 | m::Multiply( |
| 2250 | m::Multiply(m::Op(&a), m::Broadcast(m::ConstantScalar(&c1))), |
| 2251 | m::Broadcast(m::ConstantScalar(&c2))))) { |
| 2252 | TF_ASSIGN_OR_RETURN(auto* product_of_constants, |
| 2253 | MakeBinaryHlo(HloOpcode::kMultiply, c1, c2)); |
| 2254 | if (ShapeUtil::IsScalar(product_of_constants->shape()) && |
| 2255 | !ShapeUtil::IsScalar(multiply->shape())) { |
| 2256 | product_of_constants = |
| 2257 | computation_->AddInstruction(HloInstruction::CreateBroadcast( |
| 2258 | multiply->shape(), product_of_constants, {})); |
| 2259 | } |
| 2260 | return ReplaceWithNewInstruction( |
| 2261 | multiply, |
| 2262 | HloInstruction::CreateBinary(multiply->shape(), HloOpcode::kMultiply, a, |
| 2263 | product_of_constants)); |
| 2264 | } |
| 2265 | |
| 2266 | VLOG(10) << "trying to transform exp(LHS) * exp(RHS) => exp(LHS+RHS) " |
| 2267 | << multiply->ToString(); |
| 2268 | if (Match(multiply, m::Multiply(m::Exp(m::Op(&lhs)), m::Exp(m::Op(&rhs))))) { |
| 2269 | auto add = computation_->AddInstruction(HloInstruction::CreateBinary( |
| 2270 | multiply->shape(), HloOpcode::kAdd, lhs, rhs)); |
| 2271 | return ReplaceWithNewInstruction( |
| 2272 | multiply, |
| 2273 | HloInstruction::CreateUnary(multiply->shape(), HloOpcode::kExp, add)); |
no test coverage detected