Pass to eliminate redundant branch and overcompute

include/tvm/tir/transform.h

@@ -834,6 +834,14 @@ TVM_DLL Pass InstrumentProfileIntrinsics();
  */
 TVM_DLL Pass DefaultGPUSchedule();
 
+/*!
+ * \brief This pass analyzes primfunc & eliminates branch introdued due to layout specific padding.
+ *  It leverages from the buffer assumptions and use the information to eliminate the branch.
+ * \note This creates more opportunity to vectorize the code.
+ * \return The Pass.
+ */
+TVM_DLL Pass UseAssumeToReduceBranches();
+
 }  // namespace transform
 }  // namespace tir
 }  // namespace tvm

python/tvm/tir/transform/transform.py

@@ -1199,3 +1199,16 @@ def DefaultGPUSchedule():
     ret: tvm.transform.Pass
     """
     return _ffi_api.DefaultGPUSchedule()  # type: ignore
+
+
+def UseAssumeToReduceBranches():
+    """This pass attempts to eliminates layout specific pad branch by overcomputing the values
+    for padded region. Eliminating the branch will help to vectorize code,
+    and improve element wise ops performance.
+
+    Returns
+    -------
+    fpass : tvm.transform.Pass
+        The result pass
+    """
+    return _ffi_api.UseAssumeToReduceBranches()  # type: ignore