| 225 | } |
| 226 | |
| 227 | bool JITFusionPass::Impl::check_shape( |
| 228 | cg::OperatorNodeBase* opr, InternalGraphGenerator* ig_gen) { |
| 229 | if (!cg::is_static_var_shape(opr->output(0))) { |
| 230 | // currently we do not handle dynamic shape in JIT |
| 231 | return false; |
| 232 | } |
| 233 | if (!(m_feature_bits & JITFeatureBits::REDUCE)) { |
| 234 | // By requiring opr output shape to be the same as final output shape, |
| 235 | // we permit only one broadcast. If multiple broadcasts are fused, |
| 236 | // together, execution would be actually slower. |
| 237 | if ((m_feature_bits & JITFeatureBits::DIMSHUFFLE) && ig_gen->has_dimshuffle() && |
| 238 | ig_gen->oprs_depended_by_dimshuffe().count(opr)) { |
| 239 | return opr->output(0)->shape().eq_shape( |
| 240 | ig_gen->oprs_depended_by_dimshuffe().at(opr)->input(0)->shape()); |
| 241 | } else { |
| 242 | return opr->output(0)->shape().eq_shape(ig_gen->output()->shape()); |
| 243 | } |
| 244 | } |
| 245 | |
| 246 | bool before_reduce = false; |
| 247 | for (auto&& op_set : ig_gen->reduce_out_var_deps()) { |
| 248 | if (op_set.second.count(opr)) { |
| 249 | before_reduce = true; |
| 250 | break; |
| 251 | } |
| 252 | } |
| 253 | |
| 254 | if (opr->same_type<JITExecutor>()) { |
| 255 | auto jit = &opr->cast_final<JITExecutor>(); |
| 256 | bool jit_has_reduce = jit->has_reduce(); |
| 257 | auto jit_inp_shp = jit->broadcasted_input_shape(); |
| 258 | if (jit_has_reduce) { |
| 259 | if (before_reduce) |
| 260 | return jit_inp_shp.eq_shape(jit->output(0)->shape()) && |
| 261 | jit_inp_shp.eq_shape(ig_gen->before_reduce_shape()); |
| 262 | else { |
| 263 | bool ret = true; |
| 264 | if (ig_gen->has_reduce()) { |
| 265 | ret &= jit_inp_shp.eq_shape(ig_gen->before_reduce_shape()); |
| 266 | } |
| 267 | ret &= jit->output(0)->shape().eq_shape(ig_gen->output()->shape()); |
| 268 | return ret; |
| 269 | } |
| 270 | } |
| 271 | } |
| 272 | |
| 273 | if (opr->same_type<opr::Reduce>()) { |
| 274 | // TODO: handle reduce target shape in sub graph (especially considering |
| 275 | // placeholder has constant shape) |
| 276 | // |
| 277 | // The best way is to have a dedicated AST for the internal graph; but |
| 278 | // we want to reuse the deduplication and gradient mechanisms from the |
| 279 | // mgb cg |
| 280 | auto reduce = &opr->cast_final<opr::Reduce>(); |
| 281 | if (before_reduce) { |
| 282 | return reduce->input(0)->shape().eq_shape(ig_gen->before_reduce_shape()) && |
| 283 | reduce->output(0)->shape().eq_shape(ig_gen->before_reduce_shape()); |
| 284 | } else { |
nothing calls this directly
no test coverage detected