From 48b4eed52feaa22879518a24769ad07219d2936e Mon Sep 17 00:00:00 2001 From: Nima Poulad Date: Tue, 12 May 2026 01:04:36 +0000 Subject: [PATCH 1/9] perf: fuse adjacent OffloadedStmt range_for tasks at JIT time The Genesis hot path emits many small back-to-back range_for offloads (zero-init, integration, contact, ...). Each is its own GPU dispatch and saturates the AMDGPU command queue even when individual kernels are cheap. Add an opt-in JIT pass that merges adjacent range_for OffloadedStmts with identical launch bounds and no inter-task races. Safety: * Disjoint resources: A and B touch unrelated ndarrays / snodes / gtmp slots. * Same-thread RAW: every shared-resource access in either body has the same per-thread, non-indirect address fingerprint, so thread T touches the same byte in both bodies and no other thread races against it. Runs after FixCrossOffloadReferences so OffloadedStmt's gtmp offsets are final when the dynamic-bound matcher runs. Gated by QD_FUSE_TASKS (off by default); QD_FUSE_TASKS_RAW=0 is a kill switch for the RAW relaxation. QD_FUSE_TASKS_DIAG=1 emits per-kernel pass diagnostics. Co-authored-by: Cursor --- quadrants/ir/transforms.h | 6 + quadrants/transforms/fuse_offloaded_tasks.cpp | 1023 +++++++++++++++++ quadrants/transforms/offload.cpp | 7 + 3 files changed, 1036 insertions(+) create mode 100644 quadrants/transforms/fuse_offloaded_tasks.cpp diff --git a/quadrants/ir/transforms.h b/quadrants/ir/transforms.h index 61716689b6..0277a79509 100644 --- a/quadrants/ir/transforms.h +++ b/quadrants/ir/transforms.h @@ -117,6 +117,12 @@ bool determine_ad_stack_size(IRNode *root, const CompileConfig &config); bool constant_fold(IRNode *root); void associate_continue_scope(IRNode *root, const CompileConfig &config); void offload(IRNode *root, const CompileConfig &config); +// Conservative JIT-level kernel fusion: merges adjacent OffloadedStmt +// range_for tasks that have identical launch bounds and provably +// disjoint global access sets. Gated by the QD_FUSE_TASKS env var so +// the default code path is unchanged. See +// quadrants/transforms/fuse_offloaded_tasks.cpp for the safety model. +void fuse_offloaded_tasks(IRNode *root); bool transform_statements( IRNode *root, std::function filter, diff --git a/quadrants/transforms/fuse_offloaded_tasks.cpp b/quadrants/transforms/fuse_offloaded_tasks.cpp new file mode 100644 index 0000000000..0354690dbc --- /dev/null +++ b/quadrants/transforms/fuse_offloaded_tasks.cpp @@ -0,0 +1,1023 @@ +// SPDX-License-Identifier: MIT +// +// fuse_offloaded_tasks.cpp - JIT-level fusion of adjacent OffloadedStmt +// range_for tasks. +// +// Motivation. After irpass::offload() splits a kernel into N OffloadedStmt +// blocks, each top-level parallel range_for becomes a separate GPU kernel +// dispatch at runtime. On AMDGPU (gfx942) the command queue is saturated +// during the Genesis hot path (99.99% busy in rocprofv3 kernel-trace). +// Every extra kernel dispatch consumes queue cycles even if the kernel +// itself is cheap. Reducing dispatch count = reducing queue work = +// throughput gain. +// +// What this pass does. Walks the root block of an already-offloaded +// kernel and fuses pairs of adjacent OffloadedStmt's that have: +// 1. task_type == range_for in both +// 2. identical launch bounds (begin, end, block_dim) - either static +// (compile-time-known) or dynamic-via-gtmp (both offloads read +// their bound from the same global temporary slot) +// 3. *disjoint* global access resources (ndarrays / SNodes / gtmp) +// +// Condition (3) is a sufficient (not necessary) safety check. Pre-fusion, +// the kernel boundary acts as a global barrier - thread T1 in task A +// finishes before any thread starts task B. Post-fusion, every thread +// runs A's body then B's body, but threads are not synchronised +// cross-grid. If A writes a location that B reads (or vice versa, or +// both write the same location), threads from different workgroups can +// race. +// +// Disjoint-resources is conservative but tractable and catches the +// common Genesis case: multiple back-to-back zero-init loops that each +// touch a different ndarray (e.g. clear links_state.vel, clear +// dofs_state.acc, clear ...). Those produce one kernel dispatch each +// today; this pass fuses them. +// +// Same-thread RAW relaxation (QD_FUSE_TASKS_RAW=1). When two adjacent +// offloads share a resource we additionally check whether *every* +// access to that resource (in either body) uses the same per-thread +// address fingerprint. If the address is a function of the loop +// index (and loop-invariant values like args / read-only gtmp slots), +// thread T touches the same byte in both A and B and no other thread +// touches that byte - the cross-kernel-boundary RAW becomes an +// in-thread RAW, which is naturally ordered by program order. We +// fingerprint the address Stmt with LoopIndexStmt canonicalised across +// offloads (so A's and B's loop indices fingerprint identically) and +// fall back to pointer identity for Stmt kinds we can't reason about +// (which conservatively forces a mismatch). +// +// Sequencing. The pass *must* run after FixCrossOffloadReferences in +// irpass::offload() so that OffloadedStmt::begin_offset / end_offset +// are populated with their final gtmp byte offsets. Run earlier and +// the dynamic-bound matcher will see end_offset == 0 (the default +// value) on every dynamic offload and incorrectly fuse unrelated +// tasks, producing GPU memory faults. +// +// Gated behind the QD_FUSE_TASKS env var (off by default) so it can be +// A/B tested cleanly against the existing baseline. + +#include "quadrants/ir/ir.h" +#include "quadrants/ir/statements.h" +#include "quadrants/ir/transforms.h" +#include "quadrants/ir/visitors.h" +#include "quadrants/program/compile_config.h" +#include "quadrants/system/profiler.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace quadrants::lang { + +namespace { + +// Coarse-grained handle for a global resource that an IR statement +// touches. Two statements that produce the same Resource may alias; +// statements producing different Resources are guaranteed to be +// disjoint. +struct Resource { + enum class Kind : int { + Unknown = 0, + Ndarray = 1, + SNode = 2, + GlobalTmp = 3, + }; + Kind kind; + int ndarray_arg_id; + int ndarray_is_grad; + SNode *snode; + int64_t gtmp_offset; + + bool operator==(const Resource &o) const { + if (kind != o.kind) + return false; + switch (kind) { + case Kind::Unknown: + return false; // Unknowns alias with everything; never equal. + case Kind::Ndarray: + return ndarray_arg_id == o.ndarray_arg_id && + ndarray_is_grad == o.ndarray_is_grad; + case Kind::SNode: + return snode == o.snode; + case Kind::GlobalTmp: + return gtmp_offset == o.gtmp_offset; + } + return false; + } +}; + +struct ResourceHash { + size_t operator()(const Resource &r) const { + size_t h = static_cast(r.kind); + switch (r.kind) { + case Resource::Kind::Ndarray: + h ^= std::hash{}(r.ndarray_arg_id) + (h << 6) + (h >> 2); + h ^= std::hash{}(r.ndarray_is_grad) + (h << 6) + (h >> 2); + break; + case Resource::Kind::SNode: + h ^= std::hash{}(r.snode) + (h << 6) + (h >> 2); + break; + case Resource::Kind::GlobalTmp: + h ^= std::hash{}(r.gtmp_offset) + (h << 6) + (h >> 2); + break; + default: + break; + } + return h; + } +}; + +// Per-thread address fingerprint for an access. +// has_loop_idx : the fingerprinted expression depends on a +// LoopIndexStmt; different threads produce different +// addresses. +// is_indirect : the address dereferences another global pointer +// (loads from an ndarray / snode) as part of its +// index computation, e.g. `arr[other_arr[i], j]`. +// Such addresses are NOT provably injective in the +// loop index across threads (the indirection array +// could map different threads to the same byte), so +// even if A's and B's fingerprints match the +// same-thread RAW check has to reject - we can't +// rule out a cross-thread RAW. +struct AccessFp { + std::string fp; + bool has_loop_idx; + bool is_indirect; + bool operator==(const AccessFp &o) const { + return has_loop_idx == o.has_loop_idx && is_indirect == o.is_indirect && + fp == o.fp; + } +}; + +struct AccessFpHash { + size_t operator()(const AccessFp &a) const { + return std::hash{}(a.fp) ^ + (a.has_loop_idx ? 0x9e3779b97f4a7c15ULL : 0) ^ + (a.is_indirect ? 0xc2b2ae3d27d4eb4fULL : 0); + } +}; + +// Per-LocalLoadStmt resolution: the value most recently stored to its +// alloca at the point of the load, in body execution order. Computed +// per offload body so the fingerprinter can see through +// `LocalLoad(alloca)` to the expression that originally produced the +// value (typically a LoopIndexStmt or an arithmetic expression on +// one). +// +// We can't use a simpler "unique writer per alloca" map because +// Genesis emits an init store followed by a real store for almost +// every scalar local, e.g. +// alloca i; +// i = 0; // init +// i = LoopIdx / N; // real +// ...arr[i] = ... +// "unique writer" treats `i` as multi-writer and gives up. The right +// answer for each LocalLoad is the most recent LocalStoreStmt to that +// alloca in program order; that's what this map captures. +using LoadMap = std::unordered_map; + +// Recursively fingerprint an address-producing Stmt for the same- +// thread RAW check. The goal: two accesses fingerprint identically +// iff they touch the same byte for the same thread index in the +// fused kernel. +// +// Canonicalisation rules: +// LoopIndexStmt(loop=*, index=k) -> "L{k}" (offload identity dropped: +// after fusion the two loop indices are unified via +// replace_all_usages_with). +// ConstStmt -> "C{stringified value}" +// ArgLoadStmt -> "Arg{arg_id[0]}" (kernel-launch +// invariant, identical across A and B). +// ExternalTensorBasePtrStmt -> "ATB{arg_id[0]}{grad?}" +// GlobalTemporaryStmt(offset=O) -> "GT{O}" (same slot in both +// offloads). +// LocalLoadStmt(src=AllocaStmt) -> fingerprint of the alloca's +// unique writer, if any (otherwise pointer fallback). This is the +// escape hatch for the very common Genesis pattern: +// loop_idx -> alloca -> LocalLoad -> used as index +// Without it most physics-body accesses fingerprint as unrelated +// even though they all derive from the loop index. +// GlobalLoadStmt(src=...) -> "R({fp(src)})" (value-based; +// caller is responsible for proving the source slot isn't written +// between the two reads, which we enforce by also tracking the +// load's resource as part of CollectGlobalAccesses - if either +// body writes it, the resource conflict check will fail anyway). +// ExternalPtrStmt -> "EP({fp(base)},[i0,i1,...]{G?})" +// GlobalPtrStmt -> "GP({snode_ptr},[i0,i1,...])" +// MatrixPtrStmt -> "MP({fp(origin)},{fp(offset)})" +// BinaryOpStmt -> "B{op}({fp(lhs)},{fp(rhs)})" +// UnaryOpStmt -> "U{op}({fp(operand)})" +// anything else -> "@{type}_{pointer}" (conservative +// fallback; two semantically-equal Stmts with different pointers +// fingerprint as unequal, which forces a mismatch and rejects +// fusion - safe). +// +// depth is bounded to prevent runaway recursion on pathological IR. +// `inside_index` is true when fingerprinting an index sub-expression +// of an outer ExternalPtrStmt / GlobalPtrStmt / MatrixPtrStmt. Loads +// from another ndarray / snode inside an index are flagged via +// `out_is_indirect` so the caller can reject same-thread RAW for +// indirected accesses. +std::string fingerprint_value(Stmt *s, + int depth, + bool &out_has_loop_idx, + bool &out_is_indirect, + bool inside_index, + const LoadMap *load_map) { + if (depth > 16) { + return "?"; + } + if (!s) { + return "N"; + } + if (auto *li = s->cast()) { + out_has_loop_idx = true; + return "L" + std::to_string(li->index); + } + if (auto *c = s->cast()) { + return "C" + c->val.stringify(); + } + if (auto *al = s->cast()) { + int aid = al->arg_id.empty() ? -1 : al->arg_id[0]; + return "Arg" + std::to_string(aid); + } + if (auto *etbp = s->cast()) { + int aid = etbp->arg_id.empty() ? -1 : etbp->arg_id[0]; + return "ATB" + std::to_string(aid) + (etbp->is_grad ? "G" : ""); + } + if (auto *gt = s->cast()) { + return "GT" + std::to_string(gt->offset); + } + if (auto *ll = s->cast()) { + if (load_map) { + auto it = load_map->find(ll); + if (it != load_map->end()) { + return fingerprint_value(it->second, depth + 1, out_has_loop_idx, + out_is_indirect, inside_index, load_map); + } + } + // Could not resolve. Emit a stable token keyed on the alloca/src + // pointer; two unresolved LocalLoads from the same alloca within + // the same body fingerprint identically (still pointer-based, so + // doesn't match across A and B). + std::ostringstream oss; + if (auto *alloca = ll->src->cast()) { + oss << "LL_alloca_" << reinterpret_cast(alloca); + } else { + oss << "LL(" + << fingerprint_value(ll->src, depth + 1, out_has_loop_idx, + out_is_indirect, inside_index, load_map) + << ")"; + } + return oss.str(); + } + if (auto *gl = s->cast()) { + // A GlobalLoadStmt whose src is a GlobalTemporaryStmt is a + // launch-invariant scalar (loop-invariant across all threads); + // safe inside an address. A GlobalLoadStmt whose src is an + // ExternalPtrStmt / GlobalPtrStmt is an indirection through + // another ndarray / snode - the loaded value depends on thread + // index in a way we can't analyse, so flag the whole address as + // indirect. + if (inside_index && gl->src && + (gl->src->is() || gl->src->is() || + gl->src->is())) { + out_is_indirect = true; + } + return "R(" + fingerprint_value(gl->src, depth + 1, out_has_loop_idx, + out_is_indirect, inside_index, load_map) + + ")"; + } + if (auto *ep = s->cast()) { + std::string base_fp = + fingerprint_value(ep->base_ptr, depth + 1, out_has_loop_idx, + out_is_indirect, inside_index, load_map); + std::ostringstream oss; + oss << "EP(" << base_fp << ",["; + for (size_t i = 0; i < ep->indices.size(); i++) { + if (i > 0) + oss << ","; + // Indices recurse with inside_index=true so any nested + // ExternalPtr/GlobalPtr loads get flagged as indirect. + oss << fingerprint_value(ep->indices[i], depth + 1, out_has_loop_idx, + out_is_indirect, /*inside_index=*/true, + load_map); + } + oss << "]"; + if (ep->is_grad) + oss << "G"; + oss << ")"; + return oss.str(); + } + if (auto *gp = s->cast()) { + std::ostringstream oss; + oss << "GP(" << reinterpret_cast(gp->snode) << ",["; + for (size_t i = 0; i < gp->indices.size(); i++) { + if (i > 0) + oss << ","; + oss << fingerprint_value(gp->indices[i], depth + 1, out_has_loop_idx, + out_is_indirect, /*inside_index=*/true, + load_map); + } + oss << "])"; + return oss.str(); + } + if (auto *mp = s->cast()) { + std::string o_fp = + fingerprint_value(mp->origin, depth + 1, out_has_loop_idx, + out_is_indirect, inside_index, load_map); + std::string off_fp = + fingerprint_value(mp->offset, depth + 1, out_has_loop_idx, + out_is_indirect, /*inside_index=*/true, load_map); + return "MP(" + o_fp + "," + off_fp + ")"; + } + if (auto *bin = s->cast()) { + std::string lhs_fp = + fingerprint_value(bin->lhs, depth + 1, out_has_loop_idx, + out_is_indirect, inside_index, load_map); + std::string rhs_fp = + fingerprint_value(bin->rhs, depth + 1, out_has_loop_idx, + out_is_indirect, inside_index, load_map); + return "B" + std::to_string(static_cast(bin->op_type)) + "(" + + lhs_fp + "," + rhs_fp + ")"; + } + if (auto *un = s->cast()) { + std::string o_fp = + fingerprint_value(un->operand, depth + 1, out_has_loop_idx, + out_is_indirect, inside_index, load_map); + return "U" + std::to_string(static_cast(un->op_type)) + "(" + o_fp + + ")"; + } + // Unknown / unhandled node kind: fall back to pointer identity. This + // is conservative: identical-by-value Stmts with different pointers + // fingerprint as unequal, so they don't match across A and B, so the + // same-thread RAW check rejects - safe. + std::ostringstream oss; + oss << "@" << typeid(*s).name() << "_" << reinterpret_cast(s); + return oss.str(); +} + +AccessFp make_access_fp(Stmt *ptr, const LoadMap *load_map) { + AccessFp afp; + afp.has_loop_idx = false; + afp.is_indirect = false; + afp.fp = fingerprint_value(ptr, 0, afp.has_loop_idx, afp.is_indirect, + /*inside_index=*/false, load_map); + return afp; +} + +// Collects the set of AllocaStmts that are written (LocalStore or +// AtomicOp on the alloca itself) inside a block subtree. Used by the +// LocalLoadResolver below to know which alloca slots become +// unresolved after a control-flow merge. +class AllocaWriteCollector : public BasicStmtVisitor { + public: + using BasicStmtVisitor::visit; + std::unordered_set written; + + void visit(LocalStoreStmt *s) override { + if (auto *a = s->dest->cast()) + written.insert(a); + } + void visit(AtomicOpStmt *s) override { + if (auto *a = s->dest ? s->dest->cast() : nullptr) + written.insert(a); + } +}; + +// Pre-pass that walks an offload body in execution order, tracking +// the most recent LocalStoreStmt to each AllocaStmt at each program +// point. When it sees a LocalLoadStmt, it records the stored value +// that the load would observe. +// +// IfStmt / WhileStmt handling is conservative: any alloca whose value +// might differ between branches (or has been written inside a maybe- +// run loop) becomes unresolved after the construct. We don't try to +// merge branch states - a load with potentially-different values +// across paths can't be safely fingerprinted as any single one. +class LocalLoadResolver : public BasicStmtVisitor { + public: + using BasicStmtVisitor::visit; + LoadMap resolved; + + // alloca -> most recent stored value at current visit position. + // Allocas whose current value is unknown (e.g. after a conditional + // store) are absent from this map. + std::unordered_map cur; + + void visit(LocalStoreStmt *s) override { + if (auto *alloca = s->dest->cast()) { + cur[alloca] = s->val; + } + // Stores via MatrixPtrStmt / GetElementStmt etc do not update + // `cur`; the alloca's current value stays whatever it was, which + // for matrix slot stores is "unknown" - that's the desired + // behaviour. + } + + void visit(AtomicOpStmt *s) override { + if (auto *alloca = s->dest ? s->dest->cast() : nullptr) { + cur.erase(alloca); + } + } + + void visit(LocalLoadStmt *s) override { + if (auto *alloca = s->src->cast()) { + auto it = cur.find(alloca); + if (it != cur.end()) { + resolved[s] = it->second; + } + } + } + + // Drop from `cur` any alloca that is written somewhere under + // `block`, so subsequent loads correctly see "unknown". + void invalidate_writes_in(Block *block) { + if (!block) + return; + AllocaWriteCollector wc; + block->accept(&wc); + for (auto *a : wc.written) + cur.erase(a); + } + + void visit(IfStmt *if_stmt) override { + preprocess_container_stmt(if_stmt); + auto saved = cur; + if (if_stmt->true_statements) + if_stmt->true_statements->accept(this); + cur = saved; + if (if_stmt->false_statements) + if_stmt->false_statements->accept(this); + cur = saved; + invalidate_writes_in(if_stmt->true_statements.get()); + invalidate_writes_in(if_stmt->false_statements.get()); + } + + void visit(WhileStmt *stmt) override { + preprocess_container_stmt(stmt); + auto saved = cur; + if (stmt->body) + stmt->body->accept(this); + cur = saved; + invalidate_writes_in(stmt->body.get()); + } +}; + +// Walks a pointer-producing chain (MatrixPtrStmt -> ExternalPtrStmt -> +// ArgLoadStmt / ExternalTensorBasePtrStmt, or GlobalPtrStmt for SNode +// accesses) and returns the Resource it ultimately addresses, or kind = +// Unknown if we can't classify it. +Resource extract_resource(Stmt *ptr) { + Resource r; + r.kind = Resource::Kind::Unknown; + r.ndarray_arg_id = -1; + r.ndarray_is_grad = 0; + r.snode = nullptr; + r.gtmp_offset = 0; + if (!ptr) + return r; + + // Drill through MatrixPtrStmt (used for matrix/vector element access). + while (auto *mp = ptr->cast()) { + if (!mp->origin) + return r; + ptr = mp->origin; + } + + if (auto *eps = ptr->cast()) { + Stmt *base = eps->base_ptr; + if (!base) { + return r; + } + if (auto *arg = base->cast()) { + r.kind = Resource::Kind::Ndarray; + // arg_id is a vector; use the first component as the + // identifying key. We are coarse here (we don't distinguish + // sub-arrays of the same arg), which is intentional: anything + // sharing the top-level arg_id is treated as the same resource. + r.ndarray_arg_id = arg->arg_id.empty() ? -1 : arg->arg_id[0]; + r.ndarray_is_grad = eps->is_grad ? 1 : 0; + return r; + } + if (auto *etbp = base->cast()) { + r.kind = Resource::Kind::Ndarray; + r.ndarray_arg_id = etbp->arg_id.empty() ? -1 : etbp->arg_id[0]; + r.ndarray_is_grad = etbp->is_grad ? 1 : 0; + return r; + } + return r; + } + + if (auto *gps = ptr->cast()) { + r.kind = Resource::Kind::SNode; + r.snode = gps->snode; + return r; + } + + if (auto *gtmp = ptr->cast()) { + r.kind = Resource::Kind::GlobalTmp; + r.gtmp_offset = static_cast(gtmp->offset); + return r; + } + + // ThreadLocalPtrStmt / BlockLocalPtrStmt -> per-thread/per-block + // scratch, never shared cross-thread, safe to ignore. + if (ptr->is() || ptr->is()) { + return r; // Kind::Unknown but caller filters these out via the + // BasicStmtVisitor (we never call extract_resource for + // these statements). + } + + return r; +} + +// Collects the set of Resources written and read inside an +// OffloadedStmt's body. has_unknown is set if we encounter a +// global-memory operation whose target we can't classify; in that case +// the caller treats the whole offload as touching "everything" and +// refuses to fuse it. +// +// For each resource we also record the set of address fingerprints +// observed at every write / read of that resource. This is consumed +// by the same-thread RAW check (QD_FUSE_TASKS_RAW=1) to prove that +// each byte touched by both A and B is only touched by the same +// thread index in both bodies. +class CollectGlobalAccesses : public BasicStmtVisitor { + public: + using FpSet = std::unordered_set; + std::unordered_map writes; + std::unordered_map reads; + bool has_unknown = false; + // Optional. When set, fingerprinting will see through + // LocalLoad(alloca) by looking up the load's resolved value + // (the most recent LocalStore in body order). + const LoadMap *load_map = nullptr; + + using BasicStmtVisitor::visit; + + bool touches_write(const Resource &r) const { + return writes.find(r) != writes.end(); + } + bool touches_read(const Resource &r) const { + return reads.find(r) != reads.end(); + } + + void record(Stmt *ptr, bool is_write, bool is_read) { + if (ptr && (ptr->is() || ptr->is())) { + return; + } + auto r = extract_resource(ptr); + if (r.kind == Resource::Kind::Unknown) { + has_unknown = true; + return; + } + AccessFp afp = make_access_fp(ptr, load_map); + if (is_write) + writes[r].insert(afp); + if (is_read) + reads[r].insert(afp); + } + + void visit(GlobalStoreStmt *stmt) override { + record(stmt->dest, /*is_write=*/true, /*is_read=*/false); + } + + void visit(GlobalLoadStmt *stmt) override { + record(stmt->src, /*is_write=*/false, /*is_read=*/true); + } + + void visit(AtomicOpStmt *stmt) override { + // Atomic op reads+writes the same destination. + record(stmt->dest, /*is_write=*/true, /*is_read=*/true); + } + + void visit(ExternalFuncCallStmt *stmt) override { + // Conservative: a foreign call can touch anything. + has_unknown = true; + } + + void visit(InternalFuncStmt *stmt) override { + // Quadrants intrinsics: most are pure (e.g. shfl, ballot, math), + // but a few have memory side effects. Be conservative. + has_unknown = true; + } + + void visit(ClearListStmt *stmt) override { + // Touches snode meta. Don't fuse across. + has_unknown = true; + } + + void visit(RandStmt *stmt) override { + // Each invocation advances the runtime random state. Reordering + // two RandStmts changes which thread gets which random number, + // which can perturb deterministic physics. Treat as unknown. + has_unknown = true; + } + + void visit(SNodeOpStmt *stmt) override { + // Allocation / activation / deactivation of snode cells touches + // snode metadata. Conservative. + has_unknown = true; + } + + void visit(AssertStmt *stmt) override { + // Asserts have observable behavior (program termination on + // failure). Don't reorder across them. + has_unknown = true; + } +}; + +// Coarse compile-config knob: returns true iff the user has explicitly +// opted in via QD_FUSE_TASKS=1 (or =on/=true). The pass is a behavioral +// change that needs validation per workload before being made default. +bool fuse_enabled() { + static const bool enabled = []() { + const char *flag = std::getenv("QD_FUSE_TASKS"); + if (!flag || flag[0] == '\0') + return false; + // Accept "1", "on", "true", "yes" (case-insensitive); reject "0", + // "off", "false", "no". + std::string s(flag); + for (auto &c : s) + c = (char)std::tolower((unsigned char)c); + if (s == "0" || s == "off" || s == "false" || s == "no") + return false; + return true; + }(); + return enabled; +} + +bool diag_enabled() { + static const bool enabled = []() { + const char *flag = std::getenv("QD_FUSE_TASKS_DIAG"); + return flag != nullptr && flag[0] != '\0' && flag[0] != '0'; + }(); + return enabled; +} + +// Treats a resource shared between A and B as safe to fuse provided +// every access (write or read) to that resource in either body uses +// the *same* per-thread, non-indirect address fingerprint. This +// unlocks the common Genesis pattern +// arr[i] = ... (offload A) +// ... = arr[i] (offload B) +// where the kernel boundary between the writer and reader doesn't +// serve any cross-thread purpose: thread T writes arr[T] and thread T +// reads arr[T], and no other thread touches arr[T] post-fusion. +// +// Defaults to on when fusion is enabled because it gave a consistent +// throughput lift in benchmarking. Set QD_FUSE_TASKS_RAW=0 to disable +// (kill switch); the resource-disjoint policy still applies. +bool raw_enabled() { + static const bool enabled = []() { + const char *flag = std::getenv("QD_FUSE_TASKS_RAW"); + if (!flag || flag[0] == '\0') + return true; + std::string s = flag; + for (auto &c : s) + c = (char)std::tolower((unsigned char)c); + if (s == "0" || s == "off" || s == "false" || s == "no") + return false; + return true; + }(); + return enabled; +} + +// True iff a and b have identical, statically-comparable launch +// dimensions. Conservative: only accept compile-time-known bounds. +// Dynamic-bound fusion is handled separately in can_fuse_with_reason +// (after we have the access analysis) so we can prove the bound +// gtmp slot isn't written by A's body. +bool same_static_launch_dims(const OffloadedStmt *a, const OffloadedStmt *b) { + if (!a->const_begin || !a->const_end) + return false; + if (!b->const_begin || !b->const_end) + return false; + if (a->begin_value != b->begin_value) + return false; + if (a->end_value != b->end_value) + return false; + if (a->block_dim != b->block_dim) + return false; + return true; +} + +// True iff a and b have identical dynamic-via-gtmp bounds (no +// end_stmt, both read the same gtmp slot for their end). Caller must +// additionally verify no statement in A writes to that gtmp slot. +bool same_gtmp_launch_dims(const OffloadedStmt *a, const OffloadedStmt *b) { + // Bounds must match in kind (both static or both gtmp-dynamic) on + // both sides. + if (a->const_begin != b->const_begin) + return false; + if (a->const_end != b->const_end) + return false; + // We only handle the dynamic-via-gtmp case here. + if (a->const_end) + return false; // Handled by same_static_launch_dims. + if (a->end_stmt != nullptr || b->end_stmt != nullptr) + return false; + if (a->const_begin) { + if (a->begin_value != b->begin_value) + return false; + } else { + if (a->begin_offset != b->begin_offset) + return false; + } + if (a->end_offset != b->end_offset) + return false; + if (a->block_dim != b->block_dim) + return false; + return true; +} + +enum class FuseReject : int { + Ok = 0, + NotRangeFor, + TlsBls, + MemAccessOpt, + GridDim, + Bounds, + UnknownAccess, + RaceWriteWrite, + RaceWriteRead, + RaceReadWrite, + DynamicBoundRace, +}; + +const char *reject_name(FuseReject r) { + switch (r) { + case FuseReject::Ok: + return "ok"; + case FuseReject::NotRangeFor: + return "not-range-for"; + case FuseReject::TlsBls: + return "tls-or-bls-prologue"; + case FuseReject::MemAccessOpt: + return "mem-access-opt"; + case FuseReject::GridDim: + return "grid-dim-mismatch"; + case FuseReject::Bounds: + return "bounds-mismatch-or-dynamic"; + case FuseReject::UnknownAccess: + return "unknown-access"; + case FuseReject::RaceWriteWrite: + return "race-write-write"; + case FuseReject::RaceWriteRead: + return "race-a-read-b-write"; + case FuseReject::RaceReadWrite: + return "race-a-write-b-read"; + case FuseReject::DynamicBoundRace: + return "dynamic-bound-write-race"; + } + return "?"; +} + +FuseReject can_fuse_with_reason(OffloadedStmt *a, OffloadedStmt *b) { + using Type = OffloadedStmt::TaskType; + if (a->task_type != Type::range_for || b->task_type != Type::range_for) + return FuseReject::NotRangeFor; + if (a->tls_prologue || a->tls_epilogue || b->tls_prologue || + b->tls_epilogue) + return FuseReject::TlsBls; + if (a->bls_prologue || a->bls_epilogue || b->bls_prologue || + b->bls_epilogue) + return FuseReject::TlsBls; + if (!a->mem_access_opt.get_all().empty() || + !b->mem_access_opt.get_all().empty()) + return FuseReject::MemAccessOpt; + if (a->grid_dim != b->grid_dim) + return FuseReject::GridDim; + const bool static_dims = same_static_launch_dims(a, b); + // Dynamic-via-gtmp matching is on by default. It is only valid + // because this pass is sequenced *after* FixCrossOffloadReferences, + // which is where OffloadedStmt::end_offset / begin_offset get + // populated with their final gtmp byte offsets. Run earlier and + // these fields are still their default value (0) for every + // dynamic-bound offload, causing the matcher to incorrectly accept + // unrelated offloads as "same bound" and produce GPU memory faults. + const bool gtmp_dims = !static_dims && same_gtmp_launch_dims(a, b); + if (!static_dims && !gtmp_dims) + return FuseReject::Bounds; + + // Pre-pass: walk each body in execution order and resolve every + // LocalLoadStmt to the most recent LocalStore's value. This lets + // the fingerprinter see through `LocalLoad(alloca)` to the + // expression that originally produced the stored value (typically + // a LoopIndexStmt or arithmetic on one). Without this every + // physics-body access fingerprints through an opaque LocalLoadStmt + // pointer and the same-thread RAW relaxation never triggers. + LocalLoadResolver a_loads, b_loads; + a->body->accept(&a_loads); + b->body->accept(&b_loads); + + CollectGlobalAccesses a_acc, b_acc; + a_acc.load_map = &a_loads.resolved; + b_acc.load_map = &b_loads.resolved; + a->body->accept(&a_acc); + b->body->accept(&b_acc); + if (a_acc.has_unknown || b_acc.has_unknown) + return FuseReject::UnknownAccess; + + // Dynamic-bound case: pre-fusion the host evaluates A's bound, runs + // A, then re-evaluates B's bound and runs B. Post-fusion the host + // evaluates the bound *once* and uses it for both bodies. If A's + // body writes to the gtmp slot from which B reads its bound, B's + // iteration count differs across pre/post fusion - unsafe. + if (gtmp_dims) { + Resource gtmp_end; + gtmp_end.kind = Resource::Kind::GlobalTmp; + gtmp_end.gtmp_offset = static_cast(a->end_offset); + gtmp_end.ndarray_arg_id = -1; + gtmp_end.ndarray_is_grad = 0; + gtmp_end.snode = nullptr; + if (a_acc.writes.count(gtmp_end)) + return FuseReject::DynamicBoundRace; + if (!a->const_begin) { + Resource gtmp_begin = gtmp_end; + gtmp_begin.gtmp_offset = static_cast(a->begin_offset); + if (a_acc.writes.count(gtmp_begin)) + return FuseReject::DynamicBoundRace; + } + } + + // Per-resource conflict check. A pair (A,B) is racy on a resource + // iff that resource is touched by both bodies *and* at least one + // body writes to it. Under the default policy we reject any such + // conflict. Under same-thread RAW relaxation (QD_FUSE_TASKS_RAW=1) + // we instead check whether every access (write or read) to the + // shared resource in either body has the same per-thread address + // fingerprint, in which case thread T touches the same byte in + // both A and B and no other thread races against it. + const bool raw = raw_enabled(); + std::unordered_set all_resources; + for (const auto &kv : a_acc.writes) + all_resources.insert(kv.first); + for (const auto &kv : a_acc.reads) + all_resources.insert(kv.first); + for (const auto &kv : b_acc.writes) + all_resources.insert(kv.first); + for (const auto &kv : b_acc.reads) + all_resources.insert(kv.first); + + for (const auto &r : all_resources) { + const bool a_w = a_acc.writes.count(r) > 0; + const bool a_r = a_acc.reads.count(r) > 0; + const bool b_w = b_acc.writes.count(r) > 0; + const bool b_r = b_acc.reads.count(r) > 0; + if (!(a_w || a_r)) + continue; // only B touches r + if (!(b_w || b_r)) + continue; // only A touches r + if (!a_w && !b_w) + continue; // both sides read; no race + + FuseReject reason; + if (a_w && b_w) + reason = FuseReject::RaceWriteWrite; + else if (a_r && b_w) + reason = FuseReject::RaceReadWrite; + else + reason = FuseReject::RaceWriteRead; + + if (!raw) { + return reason; + } + + // Collect every (write+read) address fingerprint on this + // resource across both bodies. Safe iff there is exactly one, + // and it is per-thread. + std::unordered_set all_fps; + auto add_fps = [&]( + const std::unordered_map &m) { + auto it = m.find(r); + if (it == m.end()) + return; + for (const auto &fp : it->second) + all_fps.insert(fp); + }; + add_fps(a_acc.writes); + add_fps(a_acc.reads); + add_fps(b_acc.writes); + add_fps(b_acc.reads); + bool admit = all_fps.size() == 1 && all_fps.begin()->has_loop_idx && + !all_fps.begin()->is_indirect; + if (!admit) { + return reason; + } + // Single non-indirect per-thread fingerprint shared across A and + // B: thread T's access touches the same byte in both bodies, and + // no other thread touches that byte post-fusion. Safe. + } + return FuseReject::Ok; +} + +void merge_b_into_a(OffloadedStmt *a, OffloadedStmt *b) { + // Rewire any reference to B (e.g. LoopIndexStmt's whose .loop == + // B) so they point at A. We walk *both* offload bodies because B's + // body may contain LoopIndexStmt(loop=B,...), and A's body must not + // gain dangling references to B once we drop it. Mirrors the + // pattern used in Offloader::run when a RangeForStmt is replaced by + // its containing OffloadedStmt. + irpass::replace_all_usages_with(a, b, a); + irpass::replace_all_usages_with(b, b, a); + + // Move B's body statements into A's body. Each statement's parent + // block pointer is updated by Block::insert. + for (size_t j = 0; j < b->body->statements.size(); j++) { + a->body->insert(std::move(b->body->statements[j])); + } + b->body->statements.clear(); +} + + +int fuse_pass(Block *root_block) { + int fused_count = 0; + bool changed = true; + // Track per-reason reject counts for the diagnostic dump. + std::unordered_map reject_counts; + while (changed) { + changed = false; + auto &stmts = root_block->statements; + for (size_t i = 0; i + 1 < stmts.size(); ) { + auto *a = stmts[i]->cast(); + auto *b = stmts[i + 1]->cast(); + FuseReject reason = + (a && b) ? can_fuse_with_reason(a, b) : FuseReject::NotRangeFor; + if (reason == FuseReject::Ok) { + if (diag_enabled()) { + fmt::print(stderr, + "[fuse_offloaded_tasks] fusing offloads at indices {} and {} " + "(begin={}, end={}, block_dim={}, const_begin={}, " + "const_end={}, begin_offset={}, end_offset={})\n", + i, i + 1, a->begin_value, a->end_value, a->block_dim, + a->const_begin, a->const_end, a->begin_offset, + a->end_offset); + } + merge_b_into_a(a, b); + stmts.erase(stmts.begin() + (long)i + 1); + fused_count++; + changed = true; + // Don't advance i: try fusing the new neighbor. + } else { + reject_counts[(int)reason]++; + i++; + } + } + } + if (diag_enabled() && !reject_counts.empty()) { + fmt::print(stderr, "[fuse_offloaded_tasks] reject reasons: "); + for (const auto &kv : reject_counts) { + fmt::print(stderr, "{}={} ", reject_name((FuseReject)kv.first), + kv.second); + } + fmt::print(stderr, "\n"); + } + return fused_count; +} + +} // namespace + +namespace irpass { + +void fuse_offloaded_tasks(IRNode *root) { + QD_AUTO_PROF; + if (!fuse_enabled()) + return; + auto *root_block = root->cast(); + if (!root_block) + return; + // Diagnostic: print how many top-level offloads this kernel has and + // how many adjacent (A,B) pairs we considered, so we can tell + // whether (a) we even saw candidates and (b) which fusion check + // rejected them. + int n_offloads = 0; + int n_range_for = 0; + for (auto &s : root_block->statements) { + if (auto *o = s->cast()) { + n_offloads++; + if (o->task_type == OffloadedStmt::TaskType::range_for) + n_range_for++; + } + } + int n = fuse_pass(root_block); + if (diag_enabled()) { + fmt::print(stderr, + "[fuse_offloaded_tasks] kernel: {} top-level offloads, {} " + "range_for, fused {} pair(s) (raw={})\n", + n_offloads, n_range_for, n, raw_enabled() ? "on" : "off"); + } +} + +} // namespace irpass + +} // namespace quadrants::lang diff --git a/quadrants/transforms/offload.cpp b/quadrants/transforms/offload.cpp index 600f88f521..058b327a27 100644 --- a/quadrants/transforms/offload.cpp +++ b/quadrants/transforms/offload.cpp @@ -785,6 +785,13 @@ void offload(IRNode *root, const CompileConfig &config) { FixCrossOffloadReferences::run(root, config, &local_to_global_offset, stmt_to_offloaded, &offloaded_ranges); } + // Fuse adjacent range_for OffloadedStmt's with identical launch + // bounds and disjoint global access sets. No-op unless + // QD_FUSE_TASKS is set. Runs *after* FixCrossOffloadReferences so + // that OffloadedStmt::begin_offset / end_offset are populated with + // their final gtmp byte offsets - those are what we use to test + // "same dynamic bound" across two adjacent offloads. + fuse_offloaded_tasks(root); insert_gc(root, config); associate_continue_scope(root, config); } From b84ffab21552ce0ae68bc5f90c09639f96f4dedd Mon Sep 17 00:00:00 2001 From: Nima Poulad Date: Tue, 12 May 2026 02:08:14 +0000 Subject: [PATCH 2/9] fix: skip JIT task fusion for kernels marked cuda_graph=True The cuda_graph dispatch path relies on each top-level for-loop producing its own OffloadedStmt so the graph manager can capture them as graph nodes. The fusion pass added in 48b4eed5 was collapsing those tasks and breaking test_cuda_graph_* expectations. Plumb the @qd.kernel(cuda_graph=True) annotation onto the C++ Kernel during materialize() so compile-time passes can see it, then have fuse_offloaded_tasks return early when the owning kernel opted into cuda_graph. Also relax an incidental _num_offloaded_tasks() >= 2 assertion in test_no_cuda_graph_annotation: the test's docstring is about the graph dispatch path; the post-JIT task count for two disjoint loops is no longer part of its contract. Co-authored-by: Cursor --- python/quadrants/lang/kernel.py | 1 + quadrants/program/kernel.h | 9 +++++++++ quadrants/python/export_lang.cpp | 1 + quadrants/transforms/fuse_offloaded_tasks.cpp | 15 +++++++++++++++ tests/python/test_cuda_graph.py | 5 ++++- 5 files changed, 30 insertions(+), 1 deletion(-) diff --git a/python/quadrants/lang/kernel.py b/python/quadrants/lang/kernel.py index 46370132ce..019e92b79e 100644 --- a/python/quadrants/lang/kernel.py +++ b/python/quadrants/lang/kernel.py @@ -409,6 +409,7 @@ def materialize(self, key: "CompiledKernelKeyType | None", py_args: tuple[Any, . ) if self.fn_attrs: quadrants_kernel.set_fn_attrs(self.fn_attrs) + quadrants_kernel.use_cuda_graph = self.use_cuda_graph if _pass == 1: assert key not in self.materialized_kernels self.materialized_kernels[key] = quadrants_kernel diff --git a/quadrants/program/kernel.h b/quadrants/program/kernel.h index 9b2e3c43e4..f4fdc125fb 100644 --- a/quadrants/program/kernel.h +++ b/quadrants/program/kernel.h @@ -30,6 +30,15 @@ class QD_DLL_EXPORT Kernel : public Callable { bool is_accessor{false}; + // Mirror of the Python-side `@qd.kernel(cuda_graph=True)` annotation. + // Set by Python before compilation so IR passes can opt out of + // transformations that conflict with the cuda_graph dispatch model + // (it relies on each top-level for-loop producing its own + // OffloadedStmt). The runtime still reads the per-launch flag off + // the LaunchContextBuilder; this field is only for compile-time + // passes. + bool use_cuda_graph{false}; + Kernel(Program &program, const std::function &func, const std::string &name = "", diff --git a/quadrants/python/export_lang.cpp b/quadrants/python/export_lang.cpp index ab746f602b..de4f7b0dcd 100644 --- a/quadrants/python/export_lang.cpp +++ b/quadrants/python/export_lang.cpp @@ -663,6 +663,7 @@ void export_lang(py::module &m) { } self->fn_attrs = fn_attrs; }) + .def_readwrite("use_cuda_graph", &Kernel::use_cuda_graph) .def("insert_scalar_param", &Kernel::insert_scalar_param) .def("insert_arr_param", &Kernel::insert_arr_param) .def("insert_ndarray_param", &Kernel::insert_ndarray_param) diff --git a/quadrants/transforms/fuse_offloaded_tasks.cpp b/quadrants/transforms/fuse_offloaded_tasks.cpp index 0354690dbc..3e13d25c3f 100644 --- a/quadrants/transforms/fuse_offloaded_tasks.cpp +++ b/quadrants/transforms/fuse_offloaded_tasks.cpp @@ -61,6 +61,7 @@ #include "quadrants/ir/transforms.h" #include "quadrants/ir/visitors.h" #include "quadrants/program/compile_config.h" +#include "quadrants/program/kernel.h" #include "quadrants/system/profiler.h" #include @@ -996,6 +997,20 @@ void fuse_offloaded_tasks(IRNode *root) { auto *root_block = root->cast(); if (!root_block) return; + // Opt out of fusion entirely when the user explicitly asked for the + // cuda_graph dispatch model. The cuda_graph manager relies on each + // top-level for-loop producing its own OffloadedStmt so it can + // capture them as graph nodes; collapsing them here would silently + // change the dispatch shape the user requested. The flag is + // mirrored from Python onto Kernel::use_cuda_graph at materialize + // time, so we just inspect the first offload's owning kernel. + for (auto &s : root_block->statements) { + if (auto *o = s->cast()) { + if (o->kernel_ && o->kernel_->use_cuda_graph) + return; + break; + } + } // Diagnostic: print how many top-level offloads this kernel has and // how many adjacent (A,B) pairs we considered, so we can tell // whether (a) we even saw candidates and (b) which fusion check diff --git a/tests/python/test_cuda_graph.py b/tests/python/test_cuda_graph.py index 3d4878ab3d..965519026b 100644 --- a/tests/python/test_cuda_graph.py +++ b/tests/python/test_cuda_graph.py @@ -215,7 +215,10 @@ def two_loops(x: Annotation, y: Annotation): y = tensor_type(qd.f32, (n,)) two_loops(x, y) - assert _num_offloaded_tasks() >= 2 + # Without cuda_graph=True the JIT may fuse adjacent disjoint + # range_for tasks, so the exact task count is not part of this + # test's contract. The cuda_graph-related assertions below are. + assert _num_offloaded_tasks() >= 1 assert _cuda_graph_num_nodes() == 0 assert not _cuda_graph_used() two_loops(x, y) From 8c90abb9e015de481439430cb3f7a79ba79d0d69 Mon Sep 17 00:00:00 2001 From: Nima Poulad Date: Tue, 12 May 2026 02:43:35 +0000 Subject: [PATCH 3/9] perf: enable JIT task fusion by default Flip the QD_FUSE_TASKS default from off to on so users get the fusion + same-thread RAW lift out of the box. The env var becomes a kill switch instead of an opt-in: set QD_FUSE_TASKS=0 (or off/false/no) to bypass the pass entirely. QD_FUSE_TASKS_RAW=0 remains available to keep fusion on while disabling the same-thread RAW relaxation only. Comment updates only on the gating; no behavior change for callers that were already setting QD_FUSE_TASKS=1. Co-authored-by: Cursor --- quadrants/transforms/fuse_offloaded_tasks.cpp | 17 +++++++++-------- 1 file changed, 9 insertions(+), 8 deletions(-) diff --git a/quadrants/transforms/fuse_offloaded_tasks.cpp b/quadrants/transforms/fuse_offloaded_tasks.cpp index 3e13d25c3f..da05477821 100644 --- a/quadrants/transforms/fuse_offloaded_tasks.cpp +++ b/quadrants/transforms/fuse_offloaded_tasks.cpp @@ -53,8 +53,10 @@ // value) on every dynamic offload and incorrectly fuse unrelated // tasks, producing GPU memory faults. // -// Gated behind the QD_FUSE_TASKS env var (off by default) so it can be -// A/B tested cleanly against the existing baseline. +// Enabled by default. Set QD_FUSE_TASKS=0 (or off/false/no) to bypass +// the pass entirely. Set QD_FUSE_TASKS_RAW=0 to keep fusion on but +// disable the same-thread RAW relaxation, falling back to the +// disjoint-resources policy only. #include "quadrants/ir/ir.h" #include "quadrants/ir/statements.h" @@ -636,16 +638,15 @@ class CollectGlobalAccesses : public BasicStmtVisitor { } }; -// Coarse compile-config knob: returns true iff the user has explicitly -// opted in via QD_FUSE_TASKS=1 (or =on/=true). The pass is a behavioral -// change that needs validation per workload before being made default. +// Coarse compile-config knob: enabled by default. Set QD_FUSE_TASKS=0 +// (or off/false/no) to bypass the pass entirely; this is the kill +// switch if a workload exposes a fusion-related issue we haven't +// already covered by the safety analysis. bool fuse_enabled() { static const bool enabled = []() { const char *flag = std::getenv("QD_FUSE_TASKS"); if (!flag || flag[0] == '\0') - return false; - // Accept "1", "on", "true", "yes" (case-insensitive); reject "0", - // "off", "false", "no". + return true; std::string s(flag); for (auto &c : s) c = (char)std::tolower((unsigned char)c); From 76343d18cf87bf813a92ab01f3b1847eaff2b180 Mon Sep 17 00:00:00 2001 From: Nima Poulad Date: Tue, 12 May 2026 03:14:07 +0000 Subject: [PATCH 4/9] style: clang-format fuse_offloaded_tasks.cpp Apply pre-commit clang-format (mirrors-clang-format v19.1.7) to fix CI lint failure. Whitespace-only. Co-authored-by: Cursor --- quadrants/transforms/fuse_offloaded_tasks.cpp | 51 +++++++++---------- 1 file changed, 25 insertions(+), 26 deletions(-) diff --git a/quadrants/transforms/fuse_offloaded_tasks.cpp b/quadrants/transforms/fuse_offloaded_tasks.cpp index da05477821..d118a5157b 100644 --- a/quadrants/transforms/fuse_offloaded_tasks.cpp +++ b/quadrants/transforms/fuse_offloaded_tasks.cpp @@ -294,8 +294,9 @@ std::string fingerprint_value(Stmt *s, gl->src->is())) { out_is_indirect = true; } - return "R(" + fingerprint_value(gl->src, depth + 1, out_has_loop_idx, - out_is_indirect, inside_index, load_map) + + return "R(" + + fingerprint_value(gl->src, depth + 1, out_has_loop_idx, + out_is_indirect, inside_index, load_map) + ")"; } if (auto *ep = s->cast()) { @@ -348,8 +349,8 @@ std::string fingerprint_value(Stmt *s, std::string rhs_fp = fingerprint_value(bin->rhs, depth + 1, out_has_loop_idx, out_is_indirect, inside_index, load_map); - return "B" + std::to_string(static_cast(bin->op_type)) + "(" + - lhs_fp + "," + rhs_fp + ")"; + return "B" + std::to_string(static_cast(bin->op_type)) + "(" + lhs_fp + + "," + rhs_fp + ")"; } if (auto *un = s->cast()) { std::string o_fp = @@ -574,7 +575,8 @@ class CollectGlobalAccesses : public BasicStmtVisitor { } void record(Stmt *ptr, bool is_write, bool is_read) { - if (ptr && (ptr->is() || ptr->is())) { + if (ptr && + (ptr->is() || ptr->is())) { return; } auto r = extract_resource(ptr); @@ -787,11 +789,9 @@ FuseReject can_fuse_with_reason(OffloadedStmt *a, OffloadedStmt *b) { using Type = OffloadedStmt::TaskType; if (a->task_type != Type::range_for || b->task_type != Type::range_for) return FuseReject::NotRangeFor; - if (a->tls_prologue || a->tls_epilogue || b->tls_prologue || - b->tls_epilogue) + if (a->tls_prologue || a->tls_epilogue || b->tls_prologue || b->tls_epilogue) return FuseReject::TlsBls; - if (a->bls_prologue || a->bls_epilogue || b->bls_prologue || - b->bls_epilogue) + if (a->bls_prologue || a->bls_epilogue || b->bls_prologue || b->bls_epilogue) return FuseReject::TlsBls; if (!a->mem_access_opt.get_all().empty() || !b->mem_access_opt.get_all().empty()) @@ -898,15 +898,15 @@ FuseReject can_fuse_with_reason(OffloadedStmt *a, OffloadedStmt *b) { // resource across both bodies. Safe iff there is exactly one, // and it is per-thread. std::unordered_set all_fps; - auto add_fps = [&]( - const std::unordered_map &m) { - auto it = m.find(r); - if (it == m.end()) - return; - for (const auto &fp : it->second) - all_fps.insert(fp); - }; + auto add_fps = + [&](const std::unordered_map &m) { + auto it = m.find(r); + if (it == m.end()) + return; + for (const auto &fp : it->second) + all_fps.insert(fp); + }; add_fps(a_acc.writes); add_fps(a_acc.reads); add_fps(b_acc.writes); @@ -941,7 +941,6 @@ void merge_b_into_a(OffloadedStmt *a, OffloadedStmt *b) { b->body->statements.clear(); } - int fuse_pass(Block *root_block) { int fused_count = 0; bool changed = true; @@ -950,20 +949,20 @@ int fuse_pass(Block *root_block) { while (changed) { changed = false; auto &stmts = root_block->statements; - for (size_t i = 0; i + 1 < stmts.size(); ) { + for (size_t i = 0; i + 1 < stmts.size();) { auto *a = stmts[i]->cast(); auto *b = stmts[i + 1]->cast(); FuseReject reason = (a && b) ? can_fuse_with_reason(a, b) : FuseReject::NotRangeFor; if (reason == FuseReject::Ok) { if (diag_enabled()) { - fmt::print(stderr, + fmt::print( + stderr, "[fuse_offloaded_tasks] fusing offloads at indices {} and {} " "(begin={}, end={}, block_dim={}, const_begin={}, " "const_end={}, begin_offset={}, end_offset={})\n", i, i + 1, a->begin_value, a->end_value, a->block_dim, - a->const_begin, a->const_end, a->begin_offset, - a->end_offset); + a->const_begin, a->const_end, a->begin_offset, a->end_offset); } merge_b_into_a(a, b); stmts.erase(stmts.begin() + (long)i + 1); @@ -1028,9 +1027,9 @@ void fuse_offloaded_tasks(IRNode *root) { int n = fuse_pass(root_block); if (diag_enabled()) { fmt::print(stderr, - "[fuse_offloaded_tasks] kernel: {} top-level offloads, {} " - "range_for, fused {} pair(s) (raw={})\n", - n_offloads, n_range_for, n, raw_enabled() ? "on" : "off"); + "[fuse_offloaded_tasks] kernel: {} top-level offloads, {} " + "range_for, fused {} pair(s) (raw={})\n", + n_offloads, n_range_for, n, raw_enabled() ? "on" : "off"); } } From 5a1dcefc6f17579a12bcd7c05b2d03f6077df894 Mon Sep 17 00:00:00 2001 From: Nima Poulad Date: Wed, 13 May 2026 19:50:48 +0000 Subject: [PATCH 5/9] fix: tighten fuse_offloaded_tasks RAW relaxation to require injectivity The same-thread RAW relaxation only admitted fusion when every access to a shared resource "depended on the loop index" (AccessFp::has_loop_idx). That check was too weak: non-injective expressions like arr[i // 2], arr[i % 2], arr[min(i, K)], or arr[i & 1] syntactically contain a LoopIndexStmt but collapse multiple iterations onto the same byte. Pre-fusion the kernel boundary serialised the resulting cross-thread race; post-fusion same-thread RAW substituted each thread's local write for the cross-thread shared read, changing observable semantics. fingerprint_value now tracks injectivity per-node using local bools rather than OR-cumulating across the whole subtree, so (i + 1) / 2 is correctly classified as non-injective even though (i + 1) is. The per-operator whitelist admits LoopIndexStmt, add/sub/mul/bit_shl/bit_xor with compile-time constants (and non-zero for mul), neg, bit_not, and cast_bits; everything else (div, mod, max, min, bit_and, bit_or, bit_shr/sar, cast_value, cmp_*, math ops) is rejected. AccessFp::has_loop_idx is renamed to is_injective_in_loop_idx to reflect the stronger guarantee. Benchmarks at n_envs=8192 / 500 steps on Genesis show no regression (1.90M vs 1.79M baseline env-steps/s). Co-authored-by: Cursor --- quadrants/transforms/fuse_offloaded_tasks.cpp | 265 ++++++++++++++---- tests/python/test_fuse_offloaded_tasks.py | 259 +++++++++++++++++ 2 files changed, 476 insertions(+), 48 deletions(-) create mode 100644 tests/python/test_fuse_offloaded_tasks.py diff --git a/quadrants/transforms/fuse_offloaded_tasks.cpp b/quadrants/transforms/fuse_offloaded_tasks.cpp index d118a5157b..127c12f03a 100644 --- a/quadrants/transforms/fuse_offloaded_tasks.cpp +++ b/quadrants/transforms/fuse_offloaded_tasks.cpp @@ -36,15 +36,36 @@ // Same-thread RAW relaxation (QD_FUSE_TASKS_RAW=1). When two adjacent // offloads share a resource we additionally check whether *every* // access to that resource (in either body) uses the same per-thread -// address fingerprint. If the address is a function of the loop -// index (and loop-invariant values like args / read-only gtmp slots), -// thread T touches the same byte in both A and B and no other thread -// touches that byte - the cross-kernel-boundary RAW becomes an -// in-thread RAW, which is naturally ordered by program order. We -// fingerprint the address Stmt with LoopIndexStmt canonicalised across -// offloads (so A's and B's loop indices fingerprint identically) and -// fall back to pointer identity for Stmt kinds we can't reason about -// (which conservatively forces a mismatch). +// address fingerprint AND that address is a *provably injective* +// function of the loop index. Injectivity here means the mapping +// `i -> address(i)` is one-to-one, so thread T touches a byte that +// no other thread touches in either body - the cross-kernel-boundary +// RAW becomes an in-thread RAW which is naturally ordered by program +// order. +// +// Pure "depends on the loop index" is *not* sufficient. Non-injective +// patterns like `arr[i // 2]`, `arr[i % 2]`, `arr[min(i, K)]`, or +// `arr[i & 1]` all syntactically contain a loop index but collapse +// multiple threads onto the same byte. Pre-fusion the kernel boundary +// serialises the resulting cross-thread race; post-fusion it does not, +// so the observable result can change even though the user code was +// already racy. +// +// We therefore classify each sub-expression in the address as one of: +// - constant (Const / Arg / kernel-launch-invariant gtmp / ATB) +// - provably injective in the loop index (LoopIndexStmt, addition or +// subtraction with a compile-time constant, multiplication by a +// non-zero compile-time constant, shift-left by a compile-time +// constant, XOR with a compile-time constant, negation, bit_not, +// bitcast - all bijections on the integer lane) +// - non-injective / unclassified (everything else, including div, +// mod, max, min, bit_and, bit_or, shift-right, casts, math ops) +// An indexed access (ExternalPtr / GlobalPtr / MatrixPtr) is injective +// iff at least one of its index expressions is independently injective. +// LoopIndexStmts are canonicalised across offloads (so A's and B's loop +// indices fingerprint identically). Unhandled Stmt kinds fall back to +// pointer identity, which mismatches across A and B and forces a safe +// rejection. // // Sequencing. The pass *must* run after FixCrossOffloadReferences in // irpass::offload() so that OffloadedStmt::begin_offset / end_offset @@ -136,10 +157,34 @@ struct ResourceHash { } }; +// True iff `s` is a compile-time ConstStmt. +inline bool is_compile_time_const(Stmt *s) { + return s != nullptr && s->is(); +} + +// True iff `s` is a compile-time ConstStmt whose numeric value is non-zero. +// Used by the injectivity classifier: multiplying or shift-left'ing an +// injective expression by a *non-zero* compile-time constant preserves +// injectivity; multiplying by zero collapses every thread onto address 0. +// equal_value(static_cast(0)) constructs a TypedConstant of the +// ConstStmt's own dtype with value 0 and compares - handles all integer +// and float primitive types correctly. +inline bool is_nonzero_compile_time_const(Stmt *s) { + auto *c = s ? s->cast() : nullptr; + if (!c) + return false; + return !c->val.equal_value(static_cast(0)); +} + // Per-thread address fingerprint for an access. -// has_loop_idx : the fingerprinted expression depends on a -// LoopIndexStmt; different threads produce different -// addresses. +// is_injective_in_loop_idx : the fingerprinted expression is a provably +// injective function of the loop index - the mapping +// i -> address(i) is one-to-one, so different threads +// necessarily touch different bytes. This is strictly +// stronger than "address syntactically depends on the +// loop index": `arr[i // 2]` depends on `i` but is not +// injective in it. The same-thread RAW relaxation +// requires injectivity, not just dependence. // is_indirect : the address dereferences another global pointer // (loads from an ndarray / snode) as part of its // index computation, e.g. `arr[other_arr[i], j]`. @@ -151,18 +196,18 @@ struct ResourceHash { // rule out a cross-thread RAW. struct AccessFp { std::string fp; - bool has_loop_idx; + bool is_injective_in_loop_idx; bool is_indirect; bool operator==(const AccessFp &o) const { - return has_loop_idx == o.has_loop_idx && is_indirect == o.is_indirect && - fp == o.fp; + return is_injective_in_loop_idx == o.is_injective_in_loop_idx && + is_indirect == o.is_indirect && fp == o.fp; } }; struct AccessFpHash { size_t operator()(const AccessFp &a) const { return std::hash{}(a.fp) ^ - (a.has_loop_idx ? 0x9e3779b97f4a7c15ULL : 0) ^ + (a.is_injective_in_loop_idx ? 0x9e3779b97f4a7c15ULL : 0) ^ (a.is_indirect ? 0xc2b2ae3d27d4eb4fULL : 0); } }; @@ -222,6 +267,28 @@ using LoadMap = std::unordered_map; // fingerprint as unequal, which forces a mismatch and rejects // fusion - safe). // +// Injectivity tracking (out_is_injective_in_loop_idx): +// The caller's flag is set to true iff the expression *as a whole* is +// a provably-injective function of the loop index. Recursion uses +// *local* bools and combines them with per-operator rules: +// - LoopIndexStmt : injective. +// - add/sub i, c (c compile-time const) : injective if i is. +// - mul i, c (c non-zero const) : injective if i is. +// - bit_shl i, c (c const) : injective if i is. +// - bit_xor i, c (c const) : injective if i is. +// - neg / bit_not / cast_bits : injective if operand is. +// - EP / GP / MP : injective if at least one +// index is independently +// injective. +// - LocalLoad (resolved) : passes through to the +// stored value. +// - everything else : NOT classified injective +// (div, mod, max, min, +// bit_and, bit_or, shr, +// cmp, cast_value, math). +// We never OR-cumulate the flag across unrelated subtrees, because +// `(i + 1) / 2` is *not* injective even though `(i + 1)` is. +// // depth is bounded to prevent runaway recursion on pathological IR. // `inside_index` is true when fingerprinting an index sub-expression // of an outer ExternalPtrStmt / GlobalPtrStmt / MatrixPtrStmt. Loads @@ -230,7 +297,7 @@ using LoadMap = std::unordered_map; // indirected accesses. std::string fingerprint_value(Stmt *s, int depth, - bool &out_has_loop_idx, + bool &out_is_injective_in_loop_idx, bool &out_is_indirect, bool inside_index, const LoadMap *load_map) { @@ -240,10 +307,16 @@ std::string fingerprint_value(Stmt *s, if (!s) { return "N"; } + // LoopIndexStmt is the base case: thread T contributes exactly one + // value of L, and this is the only "interesting" injective leaf. if (auto *li = s->cast()) { - out_has_loop_idx = true; + out_is_injective_in_loop_idx = true; return "L" + std::to_string(li->index); } + // Constants and launch-invariant values. These do not depend on the + // loop index, so they never set out_is_injective_in_loop_idx. They are + // still useful "neutral" operands for injectivity-preserving ops like + // add-by-const and mul-by-const above. if (auto *c = s->cast()) { return "C" + c->val.stringify(); } @@ -258,104 +331,198 @@ std::string fingerprint_value(Stmt *s, if (auto *gt = s->cast()) { return "GT" + std::to_string(gt->offset); } + // LocalLoad of a resolved alloca is *transparent* - its injectivity + // class is exactly the resolved value's class. Pass the caller's + // flag through by reference so a loop-index-derived alloca slot can + // be classified as injective. if (auto *ll = s->cast()) { if (load_map) { auto it = load_map->find(ll); if (it != load_map->end()) { - return fingerprint_value(it->second, depth + 1, out_has_loop_idx, - out_is_indirect, inside_index, load_map); + return fingerprint_value(it->second, depth + 1, + out_is_injective_in_loop_idx, out_is_indirect, + inside_index, load_map); } } // Could not resolve. Emit a stable token keyed on the alloca/src // pointer; two unresolved LocalLoads from the same alloca within // the same body fingerprint identically (still pointer-based, so - // doesn't match across A and B). + // doesn't match across A and B). Use a throwaway local flag for + // any recursion into ll->src so an unresolved load can't leak a + // bogus injectivity claim to the caller. std::ostringstream oss; if (auto *alloca = ll->src->cast()) { oss << "LL_alloca_" << reinterpret_cast(alloca); } else { + bool ll_inj_local = false; oss << "LL(" - << fingerprint_value(ll->src, depth + 1, out_has_loop_idx, + << fingerprint_value(ll->src, depth + 1, ll_inj_local, out_is_indirect, inside_index, load_map) << ")"; } return oss.str(); } + // GlobalLoadStmt: the loaded value is a runtime quantity. Even if it + // is derived from an injective address inside its src, the loaded + // *value* is not a statically known function of the loop index and + // is not injective in it in general. Recurse into the src with a + // local flag so the loaded value's injectivity claim is discarded. + // Indirect-flag still propagates via out_is_indirect by reference. if (auto *gl = s->cast()) { - // A GlobalLoadStmt whose src is a GlobalTemporaryStmt is a - // launch-invariant scalar (loop-invariant across all threads); - // safe inside an address. A GlobalLoadStmt whose src is an - // ExternalPtrStmt / GlobalPtrStmt is an indirection through - // another ndarray / snode - the loaded value depends on thread - // index in a way we can't analyse, so flag the whole address as - // indirect. if (inside_index && gl->src && (gl->src->is() || gl->src->is() || gl->src->is())) { out_is_indirect = true; } + bool gl_inj_local = false; return "R(" + - fingerprint_value(gl->src, depth + 1, out_has_loop_idx, - out_is_indirect, inside_index, load_map) + + fingerprint_value(gl->src, depth + 1, gl_inj_local, out_is_indirect, + inside_index, load_map) + ")"; } + // ExternalPtr / GlobalPtr / MatrixPtr: the address is injective iff + // at least one index sub-expression is independently injective. We + // recurse into each index with a *local* flag and only set the + // caller's flag once at the end. The base-ptr / origin sub-tree is + // loop-invariant so its injectivity claim is discarded. if (auto *ep = s->cast()) { + bool base_inj_local = false; std::string base_fp = - fingerprint_value(ep->base_ptr, depth + 1, out_has_loop_idx, + fingerprint_value(ep->base_ptr, depth + 1, base_inj_local, out_is_indirect, inside_index, load_map); + bool any_idx_inj = false; std::ostringstream oss; oss << "EP(" << base_fp << ",["; for (size_t i = 0; i < ep->indices.size(); i++) { if (i > 0) oss << ","; - // Indices recurse with inside_index=true so any nested - // ExternalPtr/GlobalPtr loads get flagged as indirect. - oss << fingerprint_value(ep->indices[i], depth + 1, out_has_loop_idx, + bool idx_inj_local = false; + oss << fingerprint_value(ep->indices[i], depth + 1, idx_inj_local, out_is_indirect, /*inside_index=*/true, load_map); + if (idx_inj_local) + any_idx_inj = true; } oss << "]"; if (ep->is_grad) oss << "G"; oss << ")"; + if (any_idx_inj) + out_is_injective_in_loop_idx = true; return oss.str(); } if (auto *gp = s->cast()) { + bool any_idx_inj = false; std::ostringstream oss; oss << "GP(" << reinterpret_cast(gp->snode) << ",["; for (size_t i = 0; i < gp->indices.size(); i++) { if (i > 0) oss << ","; - oss << fingerprint_value(gp->indices[i], depth + 1, out_has_loop_idx, + bool idx_inj_local = false; + oss << fingerprint_value(gp->indices[i], depth + 1, idx_inj_local, out_is_indirect, /*inside_index=*/true, load_map); + if (idx_inj_local) + any_idx_inj = true; } oss << "])"; + if (any_idx_inj) + out_is_injective_in_loop_idx = true; return oss.str(); } if (auto *mp = s->cast()) { + bool origin_inj_local = false; + bool offset_inj_local = false; std::string o_fp = - fingerprint_value(mp->origin, depth + 1, out_has_loop_idx, + fingerprint_value(mp->origin, depth + 1, origin_inj_local, out_is_indirect, inside_index, load_map); std::string off_fp = - fingerprint_value(mp->offset, depth + 1, out_has_loop_idx, + fingerprint_value(mp->offset, depth + 1, offset_inj_local, out_is_indirect, /*inside_index=*/true, load_map); + if (origin_inj_local || offset_inj_local) + out_is_injective_in_loop_idx = true; return "MP(" + o_fp + "," + off_fp + ")"; } + // BinaryOp: per-operator injectivity. We recurse with local flags so + // that, e.g., (i + 1) is correctly classified as injective while + // (i + 1) // 2 - whose left operand recursively contains an + // injective sub-expression - is correctly classified as NOT + // injective. if (auto *bin = s->cast()) { + bool lhs_inj_local = false; + bool rhs_inj_local = false; std::string lhs_fp = - fingerprint_value(bin->lhs, depth + 1, out_has_loop_idx, - out_is_indirect, inside_index, load_map); + fingerprint_value(bin->lhs, depth + 1, lhs_inj_local, out_is_indirect, + inside_index, load_map); std::string rhs_fp = - fingerprint_value(bin->rhs, depth + 1, out_has_loop_idx, - out_is_indirect, inside_index, load_map); + fingerprint_value(bin->rhs, depth + 1, rhs_inj_local, out_is_indirect, + inside_index, load_map); + bool this_inj = false; + switch (bin->op_type) { + // i + c, i - c, c + i, c - i: bijection on the lane's range when + // c is a compile-time constant (which it is for stride/offset + // arithmetic emitted by the front-end). + case BinaryOpType::add: + case BinaryOpType::sub: + this_inj = (lhs_inj_local && is_compile_time_const(bin->rhs)) || + (rhs_inj_local && is_compile_time_const(bin->lhs)); + break; + // i * c, c * i: bijection iff c is a non-zero compile-time + // constant. c == 0 collapses every thread onto address 0; we + // reject that case explicitly. Overflow inside int32 is not a + // safety issue for fusion because pre- and post-fusion produce + // the same wrapped result, but a non-injective wrap-collision + // is impossible across a single loop's range for typical stride + // constants - the operands here are emitted by index arithmetic + // and would have already broken the kernel if they wrapped. + case BinaryOpType::mul: + this_inj = (lhs_inj_local && is_nonzero_compile_time_const(bin->rhs)) || + (rhs_inj_local && is_nonzero_compile_time_const(bin->lhs)); + break; + // i << c: equivalent to i * 2^c, bijective for any const c. + case BinaryOpType::bit_shl: + this_inj = lhs_inj_local && is_compile_time_const(bin->rhs); + break; + // i ^ c, c ^ i: XOR with a constant is a bijection on the + // integer bit pattern. Used by some hash-style index codes. + case BinaryOpType::bit_xor: + this_inj = (lhs_inj_local && is_compile_time_const(bin->rhs)) || + (rhs_inj_local && is_compile_time_const(bin->lhs)); + break; + // Everything else - truediv, floordiv, mod, max, min, bit_and, + // bit_or, bit_shr, bit_sar, pow, atan2, cmp_*, logical_*, ... - + // is NOT provably injective in the loop index and falls into the + // conservative "reject" bucket. This is what catches the + // arr[i // 2] / arr[i % 2] / arr[min(i, K)] / arr[i & 1] family + // of patterns the reviewer flagged. + default: + break; + } + if (this_inj) + out_is_injective_in_loop_idx = true; return "B" + std::to_string(static_cast(bin->op_type)) + "(" + lhs_fp + "," + rhs_fp + ")"; } + // UnaryOp: only the handful that are bijections preserve + // injectivity; everything else (cast_value, sqrt, abs, sgn, sin, + // cos, log, exp, ...) is rejected. if (auto *un = s->cast()) { + bool op_inj_local = false; std::string o_fp = - fingerprint_value(un->operand, depth + 1, out_has_loop_idx, - out_is_indirect, inside_index, load_map); + fingerprint_value(un->operand, depth + 1, op_inj_local, out_is_indirect, + inside_index, load_map); + bool this_inj = false; + switch (un->op_type) { + case UnaryOpType::neg: + case UnaryOpType::bit_not: + case UnaryOpType::cast_bits: + this_inj = op_inj_local; + break; + default: + break; + } + if (this_inj) + out_is_injective_in_loop_idx = true; return "U" + std::to_string(static_cast(un->op_type)) + "(" + o_fp + ")"; } @@ -370,10 +537,11 @@ std::string fingerprint_value(Stmt *s, AccessFp make_access_fp(Stmt *ptr, const LoadMap *load_map) { AccessFp afp; - afp.has_loop_idx = false; + afp.is_injective_in_loop_idx = false; afp.is_indirect = false; - afp.fp = fingerprint_value(ptr, 0, afp.has_loop_idx, afp.is_indirect, - /*inside_index=*/false, load_map); + afp.fp = + fingerprint_value(ptr, 0, afp.is_injective_in_loop_idx, afp.is_indirect, + /*inside_index=*/false, load_map); return afp; } @@ -911,7 +1079,8 @@ FuseReject can_fuse_with_reason(OffloadedStmt *a, OffloadedStmt *b) { add_fps(a_acc.reads); add_fps(b_acc.writes); add_fps(b_acc.reads); - bool admit = all_fps.size() == 1 && all_fps.begin()->has_loop_idx && + bool admit = all_fps.size() == 1 && + all_fps.begin()->is_injective_in_loop_idx && !all_fps.begin()->is_indirect; if (!admit) { return reason; diff --git a/tests/python/test_fuse_offloaded_tasks.py b/tests/python/test_fuse_offloaded_tasks.py new file mode 100644 index 0000000000..9b6d6f3880 --- /dev/null +++ b/tests/python/test_fuse_offloaded_tasks.py @@ -0,0 +1,259 @@ +"""Tests for the same-thread RAW relaxation in fuse_offloaded_tasks. + +Two adjacent offload kernels A then B that touch the same resource can +be fused into one OffloadedStmt iff every access to the resource in +either body uses the same per-thread address AND that address is a +provably *injective* function of the loop index. The injectivity check +lives in `quadrants/transforms/fuse_offloaded_tasks.cpp::fingerprint_value`. + +Two failure modes the fingerprint must guard against: + +1. Injective patterns over-rejected. arr[i], arr[2 * i + 3], arr[i ^ 1] + etc. are bijective in i, so the fused output must match the un-fused + reference byte-for-byte (every thread touches a unique byte). + +2. Non-injective patterns wrongly admitted. arr[i // 2], arr[i % 4], + arr[i & 1], arr[min(i, K)], arr[0] all collapse multiple loop + iterations onto the same byte. Pre-fusion the kernel boundary + serialises the resulting cross-thread race - all threads reading + the shared byte observe the same race-winner, so the output has a + well-defined grouping structure (e.g. out[2k] == out[2k+1] for + i // 2). Unsafe fusion via same-thread RAW relaxation would replace + the cross-thread shared read with a thread-local read of the + thread's own write, breaking the grouping invariant (each thread + sees its own value of i, so out[2k] = 2k != 2k+1 = out[2k+1]). + +The tests below assert the appropriate invariant for each pattern. +Every non-injective test below would fail before the tightened +fingerprint landed (the same-thread RAW relaxation would have admitted +the fusion, breaking the grouping invariant). +""" + +import numpy as np + +import quadrants as qd + +from tests import test_utils + +N = 1024 + + +def _run_combined(combined, arr_size): + """Run the combined kernel and return the output array as numpy.""" + a = qd.ndarray(qd.f32, shape=arr_size) + o = qd.ndarray(qd.f32, shape=N) + combined(a, o) + return o.to_numpy() + + +# --------------------------------------------------------------------------- +# Injective patterns: bijective addresses. Different threads touch +# different bytes, so the fused output is deterministic and matches +# what the per-thread arithmetic would produce. +# --------------------------------------------------------------------------- + + +@test_utils.test() +def test_fuse_raw_injective_identity(): + # arr[i] = i; out[i] = arr[i] = i + @qd.kernel + def wr(a: qd.types.NDArray, o: qd.types.NDArray): + for i in range(N): + a[i] = qd.cast(i, qd.f32) + for i in range(N): + o[i] = a[i] + + out = _run_combined(wr, N) + np.testing.assert_array_equal(out, np.arange(N, dtype=np.float32)) + + +@test_utils.test() +def test_fuse_raw_injective_const_offset(): + # arr[i + 1] = i for i in [0, N-1); out[i] = arr[i + 1] = i + @qd.kernel + def wr(a: qd.types.NDArray, o: qd.types.NDArray): + for i in range(N - 1): + a[i + 1] = qd.cast(i, qd.f32) + for i in range(N - 1): + o[i] = a[i + 1] + + out = _run_combined(wr, N) + np.testing.assert_array_equal(out[: N - 1], np.arange(N - 1, dtype=np.float32)) + + +@test_utils.test() +def test_fuse_raw_injective_const_mul(): + # arr[i * 2] = i; out[i] = arr[i * 2] = i + M = N + + @qd.kernel + def wr(a: qd.types.NDArray, o: qd.types.NDArray): + for i in range(N // 2): + a[i * 2] = qd.cast(i, qd.f32) + for i in range(N // 2): + o[i] = a[i * 2] + + out = _run_combined(wr, M) + np.testing.assert_array_equal(out[: N // 2], np.arange(N // 2, dtype=np.float32)) + + +@test_utils.test() +def test_fuse_raw_injective_affine(): + # arr[2*i + 3] = i; out[i] = arr[2*i + 3] = i + M = (2 * (N - 1) + 3) + 1 + + @qd.kernel + def wr(a: qd.types.NDArray, o: qd.types.NDArray): + for i in range(N): + a[2 * i + 3] = qd.cast(i, qd.f32) + for i in range(N): + o[i] = a[2 * i + 3] + + out = _run_combined(wr, M) + np.testing.assert_array_equal(out, np.arange(N, dtype=np.float32)) + + +@test_utils.test() +def test_fuse_raw_injective_xor_const(): + # arr[i ^ 1] = i; out[i] = arr[i ^ 1] = i (xor pairs up even/odd) + @qd.kernel + def wr(a: qd.types.NDArray, o: qd.types.NDArray): + for i in range(N): + a[i ^ 1] = qd.cast(i, qd.f32) + for i in range(N): + o[i] = a[i ^ 1] + + out = _run_combined(wr, N) + np.testing.assert_array_equal(out, np.arange(N, dtype=np.float32)) + + +@test_utils.test() +def test_fuse_raw_injective_shl_const(): + # arr[i << 1] = i; out[i] = arr[i << 1] = i + M = ((N - 1) << 1) + 1 + + @qd.kernel + def wr(a: qd.types.NDArray, o: qd.types.NDArray): + for i in range(N): + a[i << 1] = qd.cast(i, qd.f32) + for i in range(N): + o[i] = a[i << 1] + + out = _run_combined(wr, M) + np.testing.assert_array_equal(out, np.arange(N, dtype=np.float32)) + + +# --------------------------------------------------------------------------- +# Non-injective patterns: multiple loop iterations map onto the same +# byte. The kernel-boundary fence between A and B means that, after A +# completes, every thread reads the same surviving race-winner from the +# shared cell. The fused output therefore has a well-defined grouping +# structure where every group of colliding threads sees the same value. +# Same-thread RAW relaxation would substitute each thread's own write +# for the shared read, breaking the grouping. +# --------------------------------------------------------------------------- + + +@test_utils.test() +def test_fuse_raw_noninjective_floordiv(): + # arr[i // 2]: threads 2k and 2k+1 both touch arr[k]. After the + # write phase one of {2k, 2k+1} survives in arr[k]; both threads + # then read the same survivor in the read phase, so out[2k] must + # equal out[2k+1]. Unsafe fusion would give out[2k]=2k, out[2k+1]=2k+1. + @qd.kernel + def wr(a: qd.types.NDArray, o: qd.types.NDArray): + for i in range(N): + a[i // 2] = qd.cast(i, qd.f32) + for i in range(N): + o[i] = a[i // 2] + + out = _run_combined(wr, N // 2) + even = out[0:N:2] + odd = out[1:N:2] + np.testing.assert_array_equal(even, odd) + # Each surviving value must be one of {2k, 2k+1} for its slot k. + expected_lo = np.arange(0, N, 2, dtype=np.float32) + expected_hi = np.arange(1, N, 2, dtype=np.float32) + assert np.all((even == expected_lo) | (even == expected_hi)), ( + f"unexpected survivors: {even[:8].tolist()}" + ) + + +@test_utils.test() +def test_fuse_raw_noninjective_mod(): + # arr[i % 4]: 256 threads write each of arr[0..3]; all threads then + # read the survivor in their slot. Group structure: out[i] == out[i + 4]. + @qd.kernel + def wr(a: qd.types.NDArray, o: qd.types.NDArray): + for i in range(N): + a[i % 4] = qd.cast(i, qd.f32) + for i in range(N): + o[i] = a[i % 4] + + out = _run_combined(wr, 4) + for offset in range(4): + group = out[offset::4] + assert np.all(group == group[0]), ( + f"group {offset} not uniform: first 8 = {group[:8].tolist()}" + ) + + +@test_utils.test() +def test_fuse_raw_noninjective_bitand(): + # arr[i & 1]: two cells, all N threads write/read. out[i] == out[i + 2]. + @qd.kernel + def wr(a: qd.types.NDArray, o: qd.types.NDArray): + for i in range(N): + a[i & 1] = qd.cast(i, qd.f32) + for i in range(N): + o[i] = a[i & 1] + + out = _run_combined(wr, 2) + even = out[0:N:2] + odd = out[1:N:2] + assert np.all(even == even[0]), ( + f"even group not uniform: first 8 = {even[:8].tolist()}" + ) + assert np.all(odd == odd[0]), ( + f"odd group not uniform: first 8 = {odd[:8].tolist()}" + ) + + +@test_utils.test() +def test_fuse_raw_noninjective_min(): + # arr[min(i, K)]: threads i >= K all collapse to arr[K]. All + # threads with i >= K must read the same value. + K = 4 + + @qd.kernel + def wr(a: qd.types.NDArray, o: qd.types.NDArray): + for i in range(N): + a[qd.min(i, K)] = qd.cast(i, qd.f32) + for i in range(N): + o[i] = a[qd.min(i, K)] + + out = _run_combined(wr, K + 1) + clamped = out[K:] + assert np.all(clamped == clamped[0]), ( + f"clamped tail not uniform: first 8 = {clamped[:8].tolist()}" + ) + + +@test_utils.test() +def test_fuse_raw_noninjective_constant_index(): + # arr[0]: no loop-index dependence at all. All threads write arr[0] + # then all threads read arr[0]; every read must see the same + # surviving race-winner. Unsafe fusion (substituting the thread's + # own write for the shared read) would give out[i] = i, so the + # output would be non-uniform. + @qd.kernel + def wr(a: qd.types.NDArray, o: qd.types.NDArray): + for i in range(N): + a[0] = qd.cast(i, qd.f32) + for i in range(N): + o[i] = a[0] + + out = _run_combined(wr, 1) + assert np.all(out == out[0]), ( + f"out is non-uniform under arr[0]; first 8 = {out[:8].tolist()}" + ) From da36796329147eea198d5cb043670590c1dde8c3 Mon Sep 17 00:00:00 2001 From: Nima Poulad Date: Wed, 13 May 2026 20:27:00 +0000 Subject: [PATCH 6/9] style: black-format test_fuse_offloaded_tasks.py Match the pre-commit hook output: collapse the parenthesised assert message strings onto a single line. Co-authored-by: Cursor --- tests/python/test_fuse_offloaded_tasks.py | 24 ++++++----------------- 1 file changed, 6 insertions(+), 18 deletions(-) diff --git a/tests/python/test_fuse_offloaded_tasks.py b/tests/python/test_fuse_offloaded_tasks.py index 9b6d6f3880..f2290331ca 100644 --- a/tests/python/test_fuse_offloaded_tasks.py +++ b/tests/python/test_fuse_offloaded_tasks.py @@ -174,9 +174,7 @@ def wr(a: qd.types.NDArray, o: qd.types.NDArray): # Each surviving value must be one of {2k, 2k+1} for its slot k. expected_lo = np.arange(0, N, 2, dtype=np.float32) expected_hi = np.arange(1, N, 2, dtype=np.float32) - assert np.all((even == expected_lo) | (even == expected_hi)), ( - f"unexpected survivors: {even[:8].tolist()}" - ) + assert np.all((even == expected_lo) | (even == expected_hi)), f"unexpected survivors: {even[:8].tolist()}" @test_utils.test() @@ -193,9 +191,7 @@ def wr(a: qd.types.NDArray, o: qd.types.NDArray): out = _run_combined(wr, 4) for offset in range(4): group = out[offset::4] - assert np.all(group == group[0]), ( - f"group {offset} not uniform: first 8 = {group[:8].tolist()}" - ) + assert np.all(group == group[0]), f"group {offset} not uniform: first 8 = {group[:8].tolist()}" @test_utils.test() @@ -211,12 +207,8 @@ def wr(a: qd.types.NDArray, o: qd.types.NDArray): out = _run_combined(wr, 2) even = out[0:N:2] odd = out[1:N:2] - assert np.all(even == even[0]), ( - f"even group not uniform: first 8 = {even[:8].tolist()}" - ) - assert np.all(odd == odd[0]), ( - f"odd group not uniform: first 8 = {odd[:8].tolist()}" - ) + assert np.all(even == even[0]), f"even group not uniform: first 8 = {even[:8].tolist()}" + assert np.all(odd == odd[0]), f"odd group not uniform: first 8 = {odd[:8].tolist()}" @test_utils.test() @@ -234,9 +226,7 @@ def wr(a: qd.types.NDArray, o: qd.types.NDArray): out = _run_combined(wr, K + 1) clamped = out[K:] - assert np.all(clamped == clamped[0]), ( - f"clamped tail not uniform: first 8 = {clamped[:8].tolist()}" - ) + assert np.all(clamped == clamped[0]), f"clamped tail not uniform: first 8 = {clamped[:8].tolist()}" @test_utils.test() @@ -254,6 +244,4 @@ def wr(a: qd.types.NDArray, o: qd.types.NDArray): o[i] = a[0] out = _run_combined(wr, 1) - assert np.all(out == out[0]), ( - f"out is non-uniform under arr[0]; first 8 = {out[:8].tolist()}" - ) + assert np.all(out == out[0]), f"out is non-uniform under arr[0]; first 8 = {out[:8].tolist()}" From e58079467f13f3190ccac1f20a1258d713573937 Mon Sep 17 00:00:00 2001 From: Nima Poulad Date: Thu, 14 May 2026 00:04:48 +0000 Subject: [PATCH 7/9] feat: gate task fusion behind QD_AGGRESSIVE_KERNEL_FUSION env var (off by default) The pass is now opt-in: QD_FUSE_TASKS / QD_FUSE_TASKS_RAW / QD_FUSE_TASKS_DIAG are renamed to QD_AGGRESSIVE_KERNEL_FUSION{,_RAW, _DIAG}, and the master knob's default flips from on to off. With the master knob unset the offload pass returns immediately, restoring the historical no-fusion pipeline; workloads that have validated fusion under their own correctness and performance criteria can opt in by exporting QD_AGGRESSIVE_KERNEL_FUSION=1 (truthy values: 1, on, true, yes; case-insensitive). The pass currently has no register-pressure / LDS / I-cache cost model and the same-thread RAW relaxation is a syntactic safety check rather than a real loop-dependence analysis, so making it opt-in keeps the default code path unchanged for general workloads while still letting Genesis (currently the only validated user) benefit from it. A real cost heuristic and an opt-out kernel attribute can be added on top of this gate without further surgery to the default pipeline. The new test file enables the env var at module load (before any kernel compiles in the same pytest process) so the regression coverage for the injectivity check still exercises the fusion code path. Co-authored-by: Cursor --- quadrants/ir/transforms.h | 4 +- quadrants/transforms/fuse_offloaded_tasks.cpp | 43 +++++++++++-------- quadrants/transforms/offload.cpp | 2 +- tests/python/test_fuse_offloaded_tasks.py | 9 ++++ 4 files changed, 38 insertions(+), 20 deletions(-) diff --git a/quadrants/ir/transforms.h b/quadrants/ir/transforms.h index 0277a79509..814617cca6 100644 --- a/quadrants/ir/transforms.h +++ b/quadrants/ir/transforms.h @@ -119,8 +119,8 @@ void associate_continue_scope(IRNode *root, const CompileConfig &config); void offload(IRNode *root, const CompileConfig &config); // Conservative JIT-level kernel fusion: merges adjacent OffloadedStmt // range_for tasks that have identical launch bounds and provably -// disjoint global access sets. Gated by the QD_FUSE_TASKS env var so -// the default code path is unchanged. See +// disjoint global access sets. Gated by the QD_AGGRESSIVE_KERNEL_FUSION +// env var (off by default; opt-in feature). See // quadrants/transforms/fuse_offloaded_tasks.cpp for the safety model. void fuse_offloaded_tasks(IRNode *root); bool transform_statements( diff --git a/quadrants/transforms/fuse_offloaded_tasks.cpp b/quadrants/transforms/fuse_offloaded_tasks.cpp index 127c12f03a..a600aaa5ba 100644 --- a/quadrants/transforms/fuse_offloaded_tasks.cpp +++ b/quadrants/transforms/fuse_offloaded_tasks.cpp @@ -33,7 +33,7 @@ // dofs_state.acc, clear ...). Those produce one kernel dispatch each // today; this pass fuses them. // -// Same-thread RAW relaxation (QD_FUSE_TASKS_RAW=1). When two adjacent +// Same-thread RAW relaxation (QD_AGGRESSIVE_KERNEL_FUSION_RAW=1). When two adjacent // offloads share a resource we additionally check whether *every* // access to that resource (in either body) uses the same per-thread // address fingerprint AND that address is a *provably injective* @@ -74,10 +74,14 @@ // value) on every dynamic offload and incorrectly fuse unrelated // tasks, producing GPU memory faults. // -// Enabled by default. Set QD_FUSE_TASKS=0 (or off/false/no) to bypass -// the pass entirely. Set QD_FUSE_TASKS_RAW=0 to keep fusion on but -// disable the same-thread RAW relaxation, falling back to the -// disjoint-resources policy only. +// Disabled by default. Set QD_AGGRESSIVE_KERNEL_FUSION=1 (or on/true/ +// yes) to enable the pass; this is an opt-in feature - workloads must +// validate fusion against their own correctness and perf criteria +// before enabling. With the master flag on, set +// QD_AGGRESSIVE_KERNEL_FUSION_RAW=0 to keep fusion on but disable the +// same-thread RAW relaxation, falling back to the disjoint-resources +// policy only. Set QD_AGGRESSIVE_KERNEL_FUSION_DIAG=1 to print per- +// kernel fusion decisions and reject reasons to stderr. #include "quadrants/ir/ir.h" #include "quadrants/ir/statements.h" @@ -719,7 +723,7 @@ Resource extract_resource(Stmt *ptr) { // // For each resource we also record the set of address fingerprints // observed at every write / read of that resource. This is consumed -// by the same-thread RAW check (QD_FUSE_TASKS_RAW=1) to prove that +// by the same-thread RAW check (QD_AGGRESSIVE_KERNEL_FUSION_RAW=1) to prove that // each byte touched by both A and B is only touched by the same // thread index in both bodies. class CollectGlobalAccesses : public BasicStmtVisitor { @@ -808,15 +812,20 @@ class CollectGlobalAccesses : public BasicStmtVisitor { } }; -// Coarse compile-config knob: enabled by default. Set QD_FUSE_TASKS=0 -// (or off/false/no) to bypass the pass entirely; this is the kill -// switch if a workload exposes a fusion-related issue we haven't -// already covered by the safety analysis. +// Coarse compile-config knob: *disabled by default*. Set +// QD_AGGRESSIVE_KERNEL_FUSION to a truthy value (1 / on / true / yes) +// to enable the pass; default behaviour is the historical "no fusion" +// pipeline. The pass is opt-in because a) it has no register-pressure +// / LDS / I-cache cost model so it can regress occupancy-bound kernels +// after fusion, and b) the same-thread RAW relaxation is a syntactic +// safety check, not a real loop-dependence analysis. Workloads that +// have been validated under fusion (currently: Genesis on AMDGPU) can +// opt in via this env var. bool fuse_enabled() { static const bool enabled = []() { - const char *flag = std::getenv("QD_FUSE_TASKS"); + const char *flag = std::getenv("QD_AGGRESSIVE_KERNEL_FUSION"); if (!flag || flag[0] == '\0') - return true; + return false; std::string s(flag); for (auto &c : s) c = (char)std::tolower((unsigned char)c); @@ -829,7 +838,7 @@ bool fuse_enabled() { bool diag_enabled() { static const bool enabled = []() { - const char *flag = std::getenv("QD_FUSE_TASKS_DIAG"); + const char *flag = std::getenv("QD_AGGRESSIVE_KERNEL_FUSION_DIAG"); return flag != nullptr && flag[0] != '\0' && flag[0] != '0'; }(); return enabled; @@ -846,11 +855,11 @@ bool diag_enabled() { // reads arr[T], and no other thread touches arr[T] post-fusion. // // Defaults to on when fusion is enabled because it gave a consistent -// throughput lift in benchmarking. Set QD_FUSE_TASKS_RAW=0 to disable -// (kill switch); the resource-disjoint policy still applies. +// throughput lift in benchmarking. Set QD_AGGRESSIVE_KERNEL_FUSION_RAW=0 +// to disable (kill switch); the resource-disjoint policy still applies. bool raw_enabled() { static const bool enabled = []() { - const char *flag = std::getenv("QD_FUSE_TASKS_RAW"); + const char *flag = std::getenv("QD_AGGRESSIVE_KERNEL_FUSION_RAW"); if (!flag || flag[0] == '\0') return true; std::string s = flag; @@ -1022,7 +1031,7 @@ FuseReject can_fuse_with_reason(OffloadedStmt *a, OffloadedStmt *b) { // Per-resource conflict check. A pair (A,B) is racy on a resource // iff that resource is touched by both bodies *and* at least one // body writes to it. Under the default policy we reject any such - // conflict. Under same-thread RAW relaxation (QD_FUSE_TASKS_RAW=1) + // conflict. Under same-thread RAW relaxation (QD_AGGRESSIVE_KERNEL_FUSION_RAW=1) // we instead check whether every access (write or read) to the // shared resource in either body has the same per-thread address // fingerprint, in which case thread T touches the same byte in diff --git a/quadrants/transforms/offload.cpp b/quadrants/transforms/offload.cpp index 058b327a27..6168d6997a 100644 --- a/quadrants/transforms/offload.cpp +++ b/quadrants/transforms/offload.cpp @@ -787,7 +787,7 @@ void offload(IRNode *root, const CompileConfig &config) { } // Fuse adjacent range_for OffloadedStmt's with identical launch // bounds and disjoint global access sets. No-op unless - // QD_FUSE_TASKS is set. Runs *after* FixCrossOffloadReferences so + // QD_AGGRESSIVE_KERNEL_FUSION is set. Runs *after* FixCrossOffloadReferences so // that OffloadedStmt::begin_offset / end_offset are populated with // their final gtmp byte offsets - those are what we use to test // "same dynamic bound" across two adjacent offloads. diff --git a/tests/python/test_fuse_offloaded_tasks.py b/tests/python/test_fuse_offloaded_tasks.py index f2290331ca..1167d05457 100644 --- a/tests/python/test_fuse_offloaded_tasks.py +++ b/tests/python/test_fuse_offloaded_tasks.py @@ -27,8 +27,17 @@ Every non-injective test below would fail before the tightened fingerprint landed (the same-thread RAW relaxation would have admitted the fusion, breaking the grouping invariant). + +The fusion pass is opt-in (off by default), so this module enables it +explicitly via QD_AGGRESSIVE_KERNEL_FUSION=1 *before* importing +quadrants. The flag is read once on first use of the pass, so it must +be set before any kernel is compiled. """ +import os + +os.environ["QD_AGGRESSIVE_KERNEL_FUSION"] = "1" + import numpy as np import quadrants as qd From 87efa83bf75fdbc8644c3d436cc180640ab46b31 Mon Sep 17 00:00:00 2001 From: Nima Poulad Date: Thu, 14 May 2026 17:26:55 +0000 Subject: [PATCH 8/9] style: clang-format reflow comments around renamed env var The QD_AGGRESSIVE_KERNEL_FUSION* env-var names pushed three comment blocks over 80 columns; clang-format rewraps them. No code changes. Co-authored-by: Cursor --- quadrants/transforms/fuse_offloaded_tasks.cpp | 20 +++++++++---------- quadrants/transforms/offload.cpp | 8 ++++---- 2 files changed, 14 insertions(+), 14 deletions(-) diff --git a/quadrants/transforms/fuse_offloaded_tasks.cpp b/quadrants/transforms/fuse_offloaded_tasks.cpp index a600aaa5ba..e4fe3992e4 100644 --- a/quadrants/transforms/fuse_offloaded_tasks.cpp +++ b/quadrants/transforms/fuse_offloaded_tasks.cpp @@ -33,8 +33,8 @@ // dofs_state.acc, clear ...). Those produce one kernel dispatch each // today; this pass fuses them. // -// Same-thread RAW relaxation (QD_AGGRESSIVE_KERNEL_FUSION_RAW=1). When two adjacent -// offloads share a resource we additionally check whether *every* +// Same-thread RAW relaxation (QD_AGGRESSIVE_KERNEL_FUSION_RAW=1). When two +// adjacent offloads share a resource we additionally check whether *every* // access to that resource (in either body) uses the same per-thread // address fingerprint AND that address is a *provably injective* // function of the loop index. Injectivity here means the mapping @@ -723,9 +723,9 @@ Resource extract_resource(Stmt *ptr) { // // For each resource we also record the set of address fingerprints // observed at every write / read of that resource. This is consumed -// by the same-thread RAW check (QD_AGGRESSIVE_KERNEL_FUSION_RAW=1) to prove that -// each byte touched by both A and B is only touched by the same -// thread index in both bodies. +// by the same-thread RAW check (QD_AGGRESSIVE_KERNEL_FUSION_RAW=1) to prove +// that each byte touched by both A and B is only touched by the same thread +// index in both bodies. class CollectGlobalAccesses : public BasicStmtVisitor { public: using FpSet = std::unordered_set; @@ -1031,11 +1031,11 @@ FuseReject can_fuse_with_reason(OffloadedStmt *a, OffloadedStmt *b) { // Per-resource conflict check. A pair (A,B) is racy on a resource // iff that resource is touched by both bodies *and* at least one // body writes to it. Under the default policy we reject any such - // conflict. Under same-thread RAW relaxation (QD_AGGRESSIVE_KERNEL_FUSION_RAW=1) - // we instead check whether every access (write or read) to the - // shared resource in either body has the same per-thread address - // fingerprint, in which case thread T touches the same byte in - // both A and B and no other thread races against it. + // conflict. Under same-thread RAW relaxation + // (QD_AGGRESSIVE_KERNEL_FUSION_RAW=1) we instead check whether every access + // (write or read) to the shared resource in either body has the same + // per-thread address fingerprint, in which case thread T touches the same + // byte in both A and B and no other thread races against it. const bool raw = raw_enabled(); std::unordered_set all_resources; for (const auto &kv : a_acc.writes) diff --git a/quadrants/transforms/offload.cpp b/quadrants/transforms/offload.cpp index 6168d6997a..7eb99d2d1f 100644 --- a/quadrants/transforms/offload.cpp +++ b/quadrants/transforms/offload.cpp @@ -787,10 +787,10 @@ void offload(IRNode *root, const CompileConfig &config) { } // Fuse adjacent range_for OffloadedStmt's with identical launch // bounds and disjoint global access sets. No-op unless - // QD_AGGRESSIVE_KERNEL_FUSION is set. Runs *after* FixCrossOffloadReferences so - // that OffloadedStmt::begin_offset / end_offset are populated with - // their final gtmp byte offsets - those are what we use to test - // "same dynamic bound" across two adjacent offloads. + // QD_AGGRESSIVE_KERNEL_FUSION is set. Runs *after* FixCrossOffloadReferences + // so that OffloadedStmt::begin_offset / end_offset are populated with their + // final gtmp byte offsets - those are what we use to test "same dynamic + // bound" across two adjacent offloads. fuse_offloaded_tasks(root); insert_gc(root, config); associate_continue_scope(root, config); From 9ab5d1864b6cd2bdcc764c93893ee4e4475f541b Mon Sep 17 00:00:00 2001 From: Nima Poulad Date: Fri, 15 May 2026 21:14:50 +0000 Subject: [PATCH 9/9] docs: document QD_AGGRESSIVE_KERNEL_FUSION* env vars in README Adds a short "Environment variables" section listing the three knobs that gate the JIT task-fusion pass, so end users have a discoverable reference for the master flag, the same-thread RAW sub-knob, and the diagnostic toggle. Co-authored-by: Cursor --- README.md | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) diff --git a/README.md b/README.md index 5a27b61520..ec2dcd77c2 100644 --- a/README.md +++ b/README.md @@ -93,6 +93,22 @@ pip install quadrants (For how to build from source, see our CI build scripts, e.g. [linux build scripts](.github/workflows/scripts_new/linux_x86/) ) +# Environment variables + +Optional runtime knobs read once at compile time. Truthy values: `1`, `on`, `true`, `yes` (case-insensitive); anything else (including unset) is off. + +| Variable | Default | Effect | +| --- | --- | --- | +| `QD_AGGRESSIVE_KERNEL_FUSION` | off | Enables the JIT pass that fuses adjacent `range_for` `OffloadedStmt`s with matching launch bounds into a single GPU dispatch. Opt-in: workloads must validate fusion against their own correctness and perf criteria before enabling. | +| `QD_AGGRESSIVE_KERNEL_FUSION_RAW` | on (when fusion is enabled) | Allows fusion of pairs that share a resource as long as every access uses the same provably-injective per-thread address (same-thread RAW relaxation). Set to `0` to keep fusion on but restrict it to pairs with strictly disjoint resources. | +| `QD_AGGRESSIVE_KERNEL_FUSION_DIAG` | off | Prints per-kernel fusion decisions and reject reasons to stderr. Useful when debugging why a workload isn't fusing as expected. | + +Example (enable on Linux for the lifetime of one process): + +```bash +QD_AGGRESSIVE_KERNEL_FUSION=1 python my_app.py +``` + # Documentation - [docs](https://genesis-embodied-ai.github.io/quadrants/user_guide/index.html)