perf: fuse adjacent OffloadedStmt range_for tasks at JIT time

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)

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

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_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(
     IRNode *root,
     std::function<bool(Stmt *)> filter,

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<void()> &func,
          const std::string &name = "",

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)