[TIR][Transform] HoistIfThenElse added (apache#6066)

* [TIR][Transform] HoistIfThenElse added * lint error resolved * Pass position changed * pylint error resolved * CI issues resolved * Frontend tflite test case failure resolved * [1] Review comment handled * [2] Review comment handled * [3] Review comment handled * Lint error resolved
trevor-m · Aug 26, 2020 · 6d10ac9 · 6d10ac9
1 parent e68bdd0
commit 6d10ac9
Show file tree

Hide file tree

Showing 6 changed files with 652 additions and 1 deletion.
diff --git a/include/tvm/tir/transform.h b/include/tvm/tir/transform.h
@@ -338,6 +338,14 @@ TVM_DLL Pass BF16Legalize();
  */
 TVM_DLL Pass PointerValueTypeRewrite();
 
+/*!
+ * \brief Hoist loop-invariant IfThenElse nodes to
+ * outside the elligible loops.
+ *
+ * \return The pass.
+ */
+TVM_DLL Pass HoistIfThenElse();
+
 }  // namespace transform
 }  // namespace tir
 }  // namespace tvm

diff --git a/python/tvm/driver/build_module.py b/python/tvm/driver/build_module.py
@@ -179,6 +179,7 @@ def lower(sch,
         tvm.tir.transform.BF16Legalize(),
         tvm.tir.transform.NarrowDataType(32),
         tvm.tir.transform.Simplify(),
+        tvm.tir.transform.HoistIfThenElse(),
     ]
     pass_list += lower_phase1
 

diff --git a/python/tvm/tir/transform/transform.py b/python/tvm/tir/transform/transform.py
@@ -499,3 +499,12 @@ def VerifyMemory():
         The result pass
     """
     return _ffi_api.VerifyMemory()
+
+def HoistIfThenElse():
+    """Hoist loop-invariant IfThenElse nodes to outside the elligible loops.
+    Returns
+    -------
+    fpass : tvm.transform.Pass
+        The result pass
+    """
+    return _ffi_api.HoistIfThenElse()