Unverified Commit e8357626 authored by Yu Cheng's avatar Yu Cheng Committed by GitHub
Browse files

[Enhancement] Optimize loop body handling in IR (#749)



- Updated the loop body construction in `ir.cc` to conditionally include an output statement based on the analyzable condition of the `waves` variable.
- This change enhances performance by avoiding unnecessary statement wrapping when the condition is met, improving the efficiency of loop execution.
Co-authored-by: default avatarLeiWang1999 <leiwang1999@outlook.com>
parent 6b125028
...@@ -158,8 +158,13 @@ ForFrame PersistentFor(Array<PrimExpr> domain, PrimExpr wave_size, ...@@ -158,8 +158,13 @@ ForFrame PersistentFor(Array<PrimExpr> domain, PrimExpr wave_size,
tvm::tir::Call(DataType::Handle(), tvm::tl::loop_break(), {})), tvm::tir::Call(DataType::Handle(), tvm::tl::loop_break(), {})),
Stmt()); Stmt());
Stmt outer = For(loop_var, 0, waves, ForKind::kSerial, arith::Analyzer analyzer;
SeqStmt({out_if, body}), std::nullopt, anno); Stmt new_body = body;
if (analyzer.CanProveGreaterEqual(waves, 2)) {
new_body = SeqStmt({out_if, body});
}
Stmt outer =
For(loop_var, 0, waves, ForKind::kSerial, new_body, std::nullopt, anno);
for (int i = 0; i < vars.size() - 1; ++i) { for (int i = 0; i < vars.size() - 1; ++i) {
outer = tvm::tir::LetStmt(vars[i], idxs[i + 1], outer); outer = tvm::tir::LetStmt(vars[i], idxs[i + 1], outer);
} }
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment