trending/tilelang
1
0
Fork 0
mirror of https://github.com/tile-ai/tilelang.git synced 2026-10-21 06:40:23 +08:00
Code Issues Packages Projects Releases Wiki Activity

[Layout] Opaque tile-op loops must not self-plan their layout

Browse Source 
A parallel loop whose body contains an opaque tile-op call (e.g. a
reducer update) has no visible write to anchor layout inference. In
free mode such a loop used to self-plan a partition, and the
register-count-based attempt selection then preferred it: the plan
minimizes this loop's own footprint (e.g. one x-vector element per
thread) while silently forcing a column-major layout onto the fragment
operands it reads, degrading their producing copies from vectorized to
scalar accesses. On a 4096x4096 fp16 GEMV this cost 40% end-to-end
versus the legacy reducer lowering.

Fix (generic, not reducer-aware): when a loop carries an opaque
tile-op call, prefer propagating its layout from a fragment operand
that is addressed purely by the parallel loop vars, deferring while
such an operand is still unsolved so the producing ops fix the
operands first. Self-planning remains the fallback when no operand can
drive the loop (e.g. operands indexed by inner serial reduction vars,
which require the loop to own the operand layout instead) or when
operand propagation cannot express the loop's iteration space.

With the fix the GEMV benchmark matches the legacy lowering (0.075 ms
vs 0.074 ms) and narrow-plan row reductions stay ahead (AllReduce<16>
with one partial slot per thread, ~6% faster than legacy); 129 tests
pass across the reducer v1/v2 and copy suites.
This commit is contained in:
LeiWang1999
2026-08-12 02:17:23 +08:00
parent 122764554b
commit 6749631b63
2 changed files with 97 additions and 18 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/op/parallel.cc
+93 -18
View File
@@ -12,6 +12,7 @@
#include <algorithm>
#include <tvm/tirx/analysis.h>
#include <tvm/tirx/op.h>
#include <unordered_set>
#include "../layout/layout.h"
#include "../layout/utils.h"
@@ -19,6 +20,7 @@
#include "../transform/loop_vectorize.h"
#include "arith/int_operator.h"
#include "backend/common/target_utils.h"
#include "region.h"
#include "span_utils.h"
#include "utils.h"
@@ -178,7 +180,10 @@ ParallelOpNode::ParallelOpNode(For root) : root_(root), V(this) {
<< ". Ignore override.";
}
}
// Collect cross-thread access info and buffer store info.
// Collect cross-thread access info, buffer store info, and opaque tile-op
// calls (side effects the layout system cannot see, e.g. reducer updates).
static const auto op_builder_map =
Op::GetAttrMap<OpBuilderFunc>("TLOpBuilder");
PostOrderVisit(root_, [&](const ObjectRef &obj) {
if (const auto *store = obj.as<BufferStoreNode>()) {
auto buffer = store->buffer;
@@ -192,6 +197,14 @@ ParallelOpNode::ParallelOpNode(For root) : root_(root), V(this) {
if (IsSharedBuffer(load->buffer) || IsGlobalBuffer(load->buffer)) {
has_cross_thread_access_ = true;
}
} else if (const auto *call = obj.as<CallNode>()) {
if (const auto *op_node = call->op.as<OpNode>()) {
Op op_ref = tvm::ffi::GetRef<Op>(op_node);
// tl.region is a transport-only argument wrapper, not a side effect.
if (op_builder_map.count(op_ref) && !op_ref.same_as(RegionOp::Get())) {
has_opaque_tile_op_ = true;
}
}
}
});
}
@@ -483,27 +496,89 @@ LayoutMap ParallelOpNode::InferLayout(const LayoutInferArgs &layout_args,
bool selected_plan_candidate = false;
if (read_source_buffer.defined() && allow_layout_propgate) {
candidate_from_buffer =
ComputeLoopLayoutFromBuffer(read_source_buffer, layout_args);
if (has_opaque_tile_op_) {
// For opaque tile-op loops the operand is only a PREFERRED driver:
// when its indices cannot express this loop's iteration space
// (e.g. they involve inner serial vars), fall through to
// self-planning instead of failing the whole inference attempt.
try {
candidate_from_buffer =
ComputeLoopLayoutFromBuffer(read_source_buffer, layout_args);
} catch (const NormalizeIterException &e) {
DLOG(INFO) << "[FreeInfer] operand cannot drive opaque tile-op "
"loop: "
<< e.what();
}
} else {
candidate_from_buffer =
ComputeLoopLayoutFromBuffer(read_source_buffer, layout_args);
}
}
// try to infer loop layout with two mechanisms and choose the best one
{
candidate_from_plan = ComputePlanCandidate(layout_args);
// A fragment operand can DRIVE this loop's layout only when it is
// addressed purely by the parallel loop vars; indices involving other
// vars (e.g. an inner serial reduction var) require the LOOP to own the
// operand's layout instead, so such operands must not hold us back.
std::unordered_set<const VarNode *> loop_var_set;
for (const auto &iv : loop_vars_) {
loop_var_set.insert(iv->var.get());
}
bool unsolved_drivable_operand = false;
for (const auto &buffer : fragment_buffers) {
bool drivable = true;
for (const auto &index : GetAccessInfo(buffer).indices) {
if (tirx::UsesVar(index, [&](const VarNode *var) {
return !loop_var_set.count(var);
})) {
drivable = false;
break;
}
}
if (drivable && !layout_args.layout_map.count(buffer)) {
unsolved_drivable_operand = true;
break;
}
}
if (has_opaque_tile_op_ &&
(candidate_from_buffer.defined() || unsolved_drivable_operand)) {
// The loop carries an opaque side effect (e.g. a reducer update) and
// therefore has no visible write to anchor it. Self-planning such a
// loop optimizes only its own register count and can force degraded
// layouts onto the fragments it reads (their producing copies then
// lose vectorization). Propagate from a fragment operand instead;
// while a drivable operand is still unsolved, defer so the worklist
// lets the producing ops fix the operands first. Self-planning below
// remains the fallback when no operand can drive the loop, validated
// against the solved operand layouts.
if (candidate_from_buffer.defined()) {
loop_layout_ = candidate_from_buffer;
DLOG(INFO) << "[FreeInfer] opaque tile-op loop follows fragment "
"operand layout.";
} else {
DLOG(INFO) << "[FreeInfer] opaque tile-op loop deferred: drivable "
"fragment operands not solved yet.";
return {};
}
} else {
// try to infer loop layout with two mechanisms and choose the best one
{
candidate_from_plan = ComputePlanCandidate(layout_args);
}
// Choose the best candidate:
if (candidate_from_buffer.defined() && candidate_from_plan.defined()) {
loop_layout_ = ChooseBestCandidate(candidate_from_buffer,
candidate_from_plan, layout_args);
} else if (candidate_from_plan.defined()) {
loop_layout_ = candidate_from_plan;
selected_plan_candidate = true;
DLOG(INFO) << "[FreeInfer] only PlanLoopPartition available, choose it.";
} else if (candidate_from_buffer.defined()) {
loop_layout_ = candidate_from_buffer;
DLOG(INFO)
<< "[FreeInfer] only compute_from_buffer available, choose it.";
// Choose the best candidate:
if (candidate_from_buffer.defined() && candidate_from_plan.defined()) {
loop_layout_ = ChooseBestCandidate(candidate_from_buffer,
candidate_from_plan, layout_args);
} else if (candidate_from_plan.defined()) {
loop_layout_ = candidate_from_plan;
selected_plan_candidate = true;
DLOG(INFO)
<< "[FreeInfer] only PlanLoopPartition available, choose it.";
} else if (candidate_from_buffer.defined()) {
loop_layout_ = candidate_from_buffer;
DLOG(INFO)
<< "[FreeInfer] only compute_from_buffer available, choose it.";
}
}
loop_layout_requires_padding_guard_ =
selected_plan_candidate && indice_map_.empty();
src/op/parallel.h
+4
View File
@@ -206,6 +206,10 @@ private:
mutable arith::Analyzer analyzer_;
// Whether the loop body has cross-thread shared/global memory access.
bool has_cross_thread_access_ = false;
// Whether the loop body contains an opaque tile-op call (e.g. a reducer
// update): a side effect the layout system cannot see. Such loops must not
// self-plan a layout in free mode; see InferLayout.
bool has_opaque_tile_op_ = false;
// Buffers that are stored to shared/global memory in the loop body.
std::vector<Buffer> store_shared_global_buffers_;
// Fragment buffers that are stored to in the loop body.