Add notation logic

Signed-off-by: Arvind Sridhar <[email protected]>

Author: ArvindSridhar

core/lowering/lowering.cpp

@@ -25,6 +25,7 @@ void LowerBlock(torch::jit::Block* b) {
 }
 
 void LowerGraph(std::shared_ptr<torch::jit::Graph>& g, LowerInfo lower_info) {
+  passes::MarkNodesForFallback(g, false);
   passes::UnpackHardSwish(g);
   torch::jit::EliminateRedundantGuards(g);
   torch::jit::RemoveListMutation(g);
@@ -56,22 +57,27 @@ void LowerGraph(std::shared_ptr<torch::jit::Graph>& g, LowerInfo lower_info) {
   passes::AliasOperators(g);
   passes::SiluToSigmoidMultipication(g);
   torch::jit::EliminateDeadCode(g);
+  passes::MarkNodesForFallback(g, true);
   LOG_GRAPH(*g);
 }
 
-torch::jit::Module LowerModule(const torch::jit::script::Module& mod) {
-  LOG_DEBUG("Input module is being frozen by torch::jit::freeze_module");
+torch::jit::Module LowerModule(const torch::jit::Module& mod, std::string method_name, std::unordered_set<std::string> forced_fallback_modules) {
+  passes::NotateModuleForFallback(mod, "", method_name, forced_fallback_modules);
+  LOG_GRAPH("After MLF notation pass: " << *mod.get_method(method_name).graph());
   auto mod_ = torch::jit::freeze_module(mod);
+  LOG_GRAPH("After freeze: " << *mod_.get_method(method_name).graph());
   return mod_;
 }
 
 std::pair<std::shared_ptr<torch::jit::Graph>, std::vector<torch::jit::IValue>> Lower(
-    const torch::jit::script::Module& mod,
-    std::string method_name,
-    LowerInfo lower_info) {
-  auto lowered_mod = lower_info.unfreeze_module ? mod : LowerModule(mod);
+    const torch::jit::Module& mod,
+    std::string method_name, const LowerInfo& lower_info) {
+  LOG_DEBUG(lower_info);
+  LOG_GRAPH("Before lowering: " << *mod.get_method(method_name).graph());
+  std::unordered_set<std::string> forced_fallback_modules(
+      lower_info.forced_fallback_modules.begin(), lower_info.forced_fallback_modules.end());
+  auto lowered_mod = lower_info.unfreeze_module ? mod : LowerModule(mod, method_name, forced_fallback_modules);
   auto g = lowered_mod.get_method(method_name).graph();
-  LOG_GRAPH(*g);
 
   LOG_GRAPH("LibTorch Lowering");
   auto graph_and_ivalues = torch::jit::LowerGraph(*g, lowered_mod._ivalue());

core/lowering/lowering.h

@@ -15,15 +15,16 @@ struct LowerInfo {
   // Since these QDQ nodes will be identical as they share same input, one of them is eliminated due to CSE lowering
   // pass. Disable this in order to not disturb TensorRT's QAT optimizations.
   bool disable_cse = false;
-};
+  std::vector<std::string> forced_fallback_modules;
+  friend std::ostream& operator<<(std::ostream& os, const LowerInfo& l);
+}
 
 void LowerBlock(torch::jit::Block* b);
 void LowerGraph(std::shared_ptr<torch::jit::Graph>& g, LowerInfo lower_info);
-torch::jit::Module LowerModule(const torch::jit::script::Module& mod);
+torch::jit::Module LowerModule(const torch::jit::Module& mod, std::string method_name, std::unordered_set<std::string> forced_fallback_modules);
 std::pair<std::shared_ptr<torch::jit::Graph>, std::vector<torch::jit::IValue>> Lower(
-    const torch::jit::script::Module& mod,
-    std::string method_name,
-    LowerInfo lower_info);
+    const torch::jit::Module& mod,
+    std::string method_name, const LowerInfo& lower_info);
 
 } // namespace lowering
 } // namespace core