This is an automated email from the ASF dual-hosted git repository.
tqchen pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/main by this push:
new ce20316a54 [REFACTOR][TIRx] Make CPU parallel scheduling and barriers
explicit (#20549)
ce20316a54 is described below
commit ce20316a54f597c83e6e22761a1d65e18425c67b
Author: Tianqi Chen <[email protected]>
AuthorDate: Mon Oct 5 13:01:58 2026 -0400
[REFACTOR][TIRx] Make CPU parallel scheduling and barriers explicit (#20549)
CPU parallel scheduling currently uses a launch-wide stride flag, and
synchronization is attached to attribute bodies. Make scheduling local
to each parallel loop with the boolean `parallel_stride_pattern`
annotation and express team synchronization with
`T.cpu_parallel_barrier()` between loop phases.
LLVM CPU codegen reads each loop's policy directly and lowers the opaque
barrier operation using the current team. Keep the existing launch
attribute, automatic launches, team sizing, captures, and nested-loop
restrictions. Retired stride and barrier pragmas report migration
instructions, including the old stride setting's effect on later loops
in the same launch.
---
include/tvm/tirx/builtin.h | 8 ++++
python/tvm/tirx/__init__.py | 2 +-
python/tvm/tirx/op.py | 14 +++++++
python/tvm/tirx/script/ir_builder/op.py | 2 +
.../tvm/tirx/script/ir_builder/parser_protocol.py | 8 +++-
src/target/llvm/codegen_cpu.cc | 44 ++++++++++++++--------
src/target/llvm/codegen_cpu.h | 1 -
src/tirx/op/builtin.cc | 8 ++++
8 files changed, 69 insertions(+), 18 deletions(-)
diff --git a/include/tvm/tirx/builtin.h b/include/tvm/tirx/builtin.h
index 4c77ce9955..054c653f22 100644
--- a/include/tvm/tirx/builtin.h
+++ b/include/tvm/tirx/builtin.h
@@ -443,6 +443,14 @@ TVM_DLL const Op& tvm_call_cpacked_lowered();
*/
TVM_DLL const Op& tvm_storage_sync();
+/*!
+ * \brief Synchronize all workers in the current CPU parallel launch.
+ *
+ * Every worker must reach this operation. It must be outside partitioned
+ * parallel loops, whose iteration counts can differ between workers.
+ */
+TVM_DLL const Op& cpu_parallel_barrier();
+
/*!
* \brief Marker where a transform should replace generated kernel
initialization.
*/
diff --git a/python/tvm/tirx/__init__.py b/python/tvm/tirx/__init__.py
index 5af0fdab56..92954666c9 100644
--- a/python/tvm/tirx/__init__.py
+++ b/python/tvm/tirx/__init__.py
@@ -67,7 +67,7 @@ from .op import tvm_stack_alloca, tvm_stack_make_shape,
tvm_stack_make_array
from .op import handle_add_byte_offset, tvm_struct_get, tvm_struct_set
from .op import address_of, assume, undef
from .op import tvm_thread_allreduce, type_annotation, tvm_access_ptr,
ptr_byte_offset
-from .op import tvm_throw_last_error
+from .op import tvm_throw_last_error, cpu_parallel_barrier
from .op import (
tvm_load_matrix_sync,
tvm_store_matrix_sync,
diff --git a/python/tvm/tirx/op.py b/python/tvm/tirx/op.py
index 917834b750..9ca2d638b4 100644
--- a/python/tvm/tirx/op.py
+++ b/python/tvm/tirx/op.py
@@ -868,6 +868,20 @@ def tvm_storage_sync(storage_scope, is_load=False,
num_blocks=-1, *, dtype="void
return call_intrin(dtype, "tirx.tvm_storage_sync", *args)
+def cpu_parallel_barrier():
+ """Synchronize all workers in the current CPU parallel launch.
+
+ Every worker must reach this operation at the same program point. Place it
+ inside a ``pragma_parallel_launch_point`` scope, outside parallel loops,
+ whose iteration counts can differ between workers. Writes before the
barrier
+ are visible to all workers after it.
+
+ To replace ``pragma_parallel_barrier_when_finish``, place this operation
+ after the former attribute body.
+ """
+ return call_intrin("void", "tirx.cpu_parallel_barrier")
+
+
def tvm_kernel_replace_point():
"""Mark where a transform should replace generated kernel
initialization."""
return call_intrin("void", "tirx.tvm_kernel_replace_point")
diff --git a/python/tvm/tirx/script/ir_builder/op.py
b/python/tvm/tirx/script/ir_builder/op.py
index 3db9cec137..e2b5304daf 100644
--- a/python/tvm/tirx/script/ir_builder/op.py
+++ b/python/tvm/tirx/script/ir_builder/op.py
@@ -624,6 +624,7 @@ tvm_store_matrix_sync = _tir_op.tvm_store_matrix_sync
tvm_storage_sync = _tir_op.tvm_storage_sync
+cpu_parallel_barrier = _tir_op.cpu_parallel_barrier
tvm_kernel_replace_point = _tir_op.tvm_kernel_replace_point
@@ -940,6 +941,7 @@ __all__ = [
"copysign",
"cos",
"cosh",
+ "cpu_parallel_barrier",
"dp4a",
"erf",
"exp",
diff --git a/python/tvm/tirx/script/ir_builder/parser_protocol.py
b/python/tvm/tirx/script/ir_builder/parser_protocol.py
index dfdde9ff4b..aecc3980a2 100644
--- a/python/tvm/tirx/script/ir_builder/parser_protocol.py
+++ b/python/tvm/tirx/script/ir_builder/parser_protocol.py
@@ -778,7 +778,13 @@ def parallel(
The maximum value of iteration.
annotations : Dict[str, Any]
- The optional annotations of the For statement.
+ The optional annotations of the For statement. On CPU,
+ ``{"parallel_stride_pattern": True}`` assigns iterations cyclically:
+ worker ``t`` executes ``t, t + P, t + 2 * P, ...`` for ``P`` workers.
+ Absent or false selects contiguous chunks. Each loop selects its own
+ policy. When migrating ``pragma_parallel_stride_pattern``, annotate
+ every intended loop, including later loops in the same launch outside
+ the former attribute body.
step : Expr
The optional step value of iteration.
diff --git a/src/target/llvm/codegen_cpu.cc b/src/target/llvm/codegen_cpu.cc
index 89ad9299c4..7833a32215 100644
--- a/src/target/llvm/codegen_cpu.cc
+++ b/src/target/llvm/codegen_cpu.cc
@@ -1000,6 +1000,15 @@ llvm::Value* CodeGenCPU::CreateIntrinsic(const CallNode*
op) {
return CreateCallPacked(op);
} else if (op->op.same_as(tirx::builtin::tvm_call_cpacked_lowered())) {
return CreateCallPacked(op);
+ } else if (op->op.same_as(tirx::builtin::cpu_parallel_barrier())) {
+ TVM_FFI_ICHECK_EQ(args.size(), 0U);
+ TVM_FFI_ICHECK(parallel_env_.penv != nullptr)
+ << "cpu_parallel_barrier requires a parallel launch";
+ TVM_FFI_ICHECK(!parallel_env_.in_parallel_loop)
+ << "cpu_parallel_barrier must be outside parallel loops so every
worker reaches it";
+ auto callee = llvm::FunctionCallee(ftype_tvm_parallel_barrier_,
RuntimeTVMParallelBarrier());
+ return builder_->CreateCall(callee,
+ {MakeValue(parallel_env_.task_id.value()),
parallel_env_.penv});
} else if (op->op.same_as(tirx::builtin::tvm_static_handle())) {
return CreateStaticHandle();
} else if (op->op.same_as(tirx::builtin::tvm_throw_last_error())) {
@@ -1128,23 +1137,16 @@ void CodeGenCPU::Dispatch_(const AttrStmtNode* op) {
this->CreateComputeScope(op);
} else if (tirx::attr::IsPragmaKey(op->attr_key)) {
if (op->attr_key == "pragma_parallel_stride_pattern") {
- TVM_FFI_ICHECK(parallel_env_.penv != nullptr)
- << "Pragma parallel_stride_pattern only valid in parallel launch";
- parallel_env_.stride_pattern = true;
- this->Dispatch(op->body);
+ TVM_FFI_THROW(ValueError)
+ << "pragma_parallel_stride_pattern is retired; annotate every
affected parallel For "
+ << "with parallel_stride_pattern=True, including later loops in the
same launch "
+ << "outside the former attribute body";
} else if (op->attr_key == "pragma_parallel_launch_point") {
CreateParallelLaunch(op->body, 0, "pragma_parallel");
} else if (op->attr_key == "pragma_parallel_barrier_when_finish") {
- TVM_FFI_ICHECK(parallel_env_.penv != nullptr)
- << "Cannot run barrier without parallel environment";
- TVM_FFI_ICHECK(!parallel_env_.in_parallel_loop)
- << "Cannot not place within parallel loop as the workload may
differ, "
- << " place it between parallel and parallel_launch_point";
- this->Dispatch(op->body);
- auto bar_callee =
- llvm::FunctionCallee(ftype_tvm_parallel_barrier_,
RuntimeTVMParallelBarrier());
- builder_->CreateCall(bar_callee,
- {MakeValue(parallel_env_.task_id.value()),
parallel_env_.penv});
+ TVM_FFI_THROW(ValueError)
+ << "pragma_parallel_barrier_when_finish is retired; place
cpu_parallel_barrier() "
+ << "after the former attribute body inside the parallel launch";
} else if (op->attr_key == tirx::attr::pragma_import_llvm) {
const StringImmNode* value = op->value.as<StringImmNode>();
TVM_FFI_ICHECK(value != nullptr);
@@ -1184,7 +1186,19 @@ void CodeGenCPU::Dispatch_(const ForNode* op) {
<< "Nested parallel loop is not supported by threadpool, try fuse
them instead";
parallel_env_.in_parallel_loop = true;
PrimExpr end = is_zero(op->min) ? op->extent :
analyzer_->Simplify(op->min + op->extent);
- if (parallel_env_.stride_pattern) {
+ bool stride_pattern = false;
+ if (auto it = op->annotations.find("parallel_stride_pattern"); it !=
op->annotations.end()) {
+ ffi::Any annotation = (*it).second;
+ if (auto value = annotation.as<bool>()) {
+ stride_pattern = value.value();
+ } else if (auto value = annotation.try_cast<IntImm>();
+ value && value.value().ty() == PrimType::Bool()) {
+ stride_pattern = is_one(value.value());
+ } else {
+ TVM_FFI_THROW(ValueError) << "parallel_stride_pattern must be a
constant boolean";
+ }
+ }
+ if (stride_pattern) {
CreateSerialFor(MakeValue(task_id), MakeValue(end),
MakeValue(num_task), op->loop_var,
op->body);
} else {
diff --git a/src/target/llvm/codegen_cpu.h b/src/target/llvm/codegen_cpu.h
index c44d2d743b..dd9dc57436 100644
--- a/src/target/llvm/codegen_cpu.h
+++ b/src/target/llvm/codegen_cpu.h
@@ -107,7 +107,6 @@ class CodeGenCPU : public CodeGenLLVM {
struct ParallelEnv {
ffi::Optional<Var> task_id;
ffi::Optional<Var> num_task;
- bool stride_pattern{false};
bool in_parallel_loop{false};
int parallel_loop_count{0};
llvm::Value* penv{nullptr};
diff --git a/src/tirx/op/builtin.cc b/src/tirx/op/builtin.cc
index cd9d4b0931..167ac2541b 100644
--- a/src/tirx/op/builtin.cc
+++ b/src/tirx/op/builtin.cc
@@ -222,6 +222,7 @@ TVM_DEFINE_CACHED_OP_GETTER(tvm_thread_invariant,
"tirx.tvm_thread_invariant")
TVM_DEFINE_CACHED_OP_GETTER(tvm_call_packed_lowered,
"tirx.tvm_call_packed_lowered")
TVM_DEFINE_CACHED_OP_GETTER(tvm_call_cpacked_lowered,
"tirx.tvm_call_cpacked_lowered")
TVM_DEFINE_CACHED_OP_GETTER(tvm_storage_sync, "tirx.tvm_storage_sync")
+TVM_DEFINE_CACHED_OP_GETTER(cpu_parallel_barrier, "tirx.cpu_parallel_barrier")
TVM_DEFINE_CACHED_OP_GETTER(tvm_kernel_replace_point,
"tirx.tvm_kernel_replace_point")
TVM_DEFINE_CACHED_OP_GETTER(tvm_warp_shuffle, "tirx.tvm_warp_shuffle")
TVM_DEFINE_CACHED_OP_GETTER(tvm_warp_shuffle_up, "tirx.tvm_warp_shuffle_up")
@@ -529,6 +530,13 @@ TVM_FFI_STATIC_INIT_BLOCK() {
.set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
.set_attr<TCallEffectKind>("TCallEffectKind",
static_cast<int64_t>(CallEffectKind::kOpaque));
+ OpDef("tirx.cpu_parallel_barrier")
+ .signature()
+ .set_attr<TFixedReturnType>("TFixedReturnType", PrimType::Void())
+ .set_attr<TScriptPrinterName>("TScriptPrinterName",
ffi::String("tirx.cpu_parallel_barrier"))
+ .set_attr<TIRxOpCategory>("TIRxOpCategory", ffi::String("builtin"))
+ .set_attr<TCallEffectKind>("TCallEffectKind",
static_cast<int64_t>(CallEffectKind::kOpaque));
+
OpDef("tirx.tvm_kernel_replace_point")
.set_attr<TFixedReturnType>("TFixedReturnType", PrimType::Void())
.set_attr<TScriptPrinterName>("TScriptPrinterName",
