| 474 | } |
| 475 | |
| 476 | JITFusionPass::JITFusionPass( |
| 477 | bool after_grad, int jit_opt_level, const JITConfig& jit_config) |
| 478 | : m_after_grad{after_grad}, m_feature_bits{JITFeatureBits::NONE} { |
| 479 | // get default config from jit_opt_level |
| 480 | JITConfig config; |
| 481 | if (jit_opt_level == 1) { |
| 482 | config.fuse_dimshuffle = JITConfig::ON; |
| 483 | config.fuse_reduce = JITConfig::OFF; |
| 484 | } else if (jit_opt_level >= 2) { |
| 485 | config.fuse_dimshuffle = JITConfig::OFF; |
| 486 | config.fuse_reduce = JITConfig::ON; |
| 487 | } |
| 488 | |
| 489 | // overwrite default config with custom settings |
| 490 | config.update(jit_config); |
| 491 | bool fuse_dimshuffle = config.fuse_dimshuffle == JITConfig::ON; |
| 492 | bool fuse_reduce = config.fuse_reduce == JITConfig::ON; |
| 493 | |
| 494 | if (fuse_dimshuffle && fuse_reduce) { |
| 495 | mgb_assert(false, "reduce and dimshuffle can not coexist now"); |
| 496 | } |
| 497 | if (fuse_dimshuffle) { |
| 498 | m_feature_bits |= JITFeatureBits::DIMSHUFFLE; |
| 499 | } |
| 500 | if (fuse_reduce) { |
| 501 | m_feature_bits |= JITFeatureBits::REDUCE; |
| 502 | } |
| 503 | } |
| 504 | |
| 505 | const char* JITFusionPass::name() const { |
| 506 | return mgb_cstr_log("fusion_pass"); |