Enhance Loop Unswitching with CallNode Exclusion Check

- Added CallCheckerExcludingIf class to ensure function calls outside of hoisted if statements are identified, preventing potential synchronization issues during loop unswitching. - Updated loop unswitching logic to incorporate the new call checker, enhancing safety and correctness. - Integrated debug print statements in OptimizeForTarget to visualize module state before and after loop unswitching. - Disabled tilelang cache in the benchmark script for improved performance.
LeiWang1999 · May 20, 2026 · Jan 28, 2026 · Jan 28, 2026 · Jan 28, 2026 · Jan 28, 2026
commit 1f06094dcb17b6e32f4a677229017bd0e4faddf2
diff --git a/src/transform/loop_unswitching.cc b/src/transform/loop_unswitching.cc
@@ -141,6 +141,36 @@ class CallNodeChecker : public ExprVisitor {
   }
 };
 
+/*!
+ * \brief Check if a statement contains any CallNode, excluding a specific If
+ *
+ * Loop unswitching is unsafe when there are function calls OUTSIDE the
+ * hoisted if statement, because those calls (originally executed by all
+ * threads together) would be split into different code paths after
+ * unswitching, potentially breaking synchronization semantics.
+ *
+ * Calls INSIDE the if are safe because they were already conditionally
+ * executed before unswitching.
+ */
+class CallCheckerExcludingIf : public StmtExprVisitor {
+public:
+  bool has_call = false;
+  const IfThenElseNode *excluded_if = nullptr;
+
+  void VisitStmt_(const IfThenElseNode *op) final {
+    if (op == excluded_if) {
+      // Skip the interior of the excluded if statement
+      return;
+    }
+    StmtExprVisitor::VisitStmt_(op);
+  }
+
+  void VisitExpr_(const CallNode *op) final {
+    has_call = true;
+    // No need to continue once we find a call
+  }
+};
+
 /*!
  * \brief Check if condition or any Let-bound variable it uses depends on loop
  * var
@@ -297,6 +327,21 @@ class LoopUnswitcher : public StmtExprMutator {
                  op->thread_binding, op->annotations);
     }
 
+    // Check if there are any function calls OUTSIDE the hoisted if statement.
+    // Calls outside the if are executed by all threads together; unswitching
+    // would split them into different code paths, breaking synchronization.
+    // Calls inside the if are already conditionally executed, so they're safe.
+    CallCheckerExcludingIf call_checker;
+    call_checker.excluded_if = finder.found;
+    call_checker(body);
+    if (call_checker.has_call) {
+      if (body.same_as(op->body)) {
+        return ffi::GetRef<Stmt>(op);
+      }
+      return For(op->loop_var, op->min, op->extent, op->kind, body,
+                 op->thread_binding, op->annotations);
+    }
+
     // Unswitch: create two loop versions
     const IfThenElseNode *if_node = finder.found;
 

diff --git a/tilelang/engine/phase.py b/tilelang/engine/phase.py
@@ -236,7 +236,11 @@ def OptimizeForTarget(mod: IRModule, target: Target) -> IRModule:
     mod = tir.transform.Simplify()(mod)
     mod = tilelang.transform.VectorizeLoop(enable_vectorize=allow_vectorize(pass_ctx=pass_ctx))(mod)
     mod = tilelang.transform.StorageRewrite()(mod)
+    print("Before StorageRewrite")
+    print(mod)
     mod = tilelang.transform.LoopUnswitching()(mod)
+    print("After LoopUnswitching")
+    print(mod)
     mod = tilelang.transform.UnrollLoop()(mod)
     mod = tir.transform.RenormalizeSplitPattern()(mod)
     mod = tir.transform.Simplify()(mod)