src/jit/trace_plan.h
| Line | Branch | Exec | Source |
|---|---|---|---|
| 1 | #pragma once | ||
| 2 | |||
| 3 | // Device-neutral fusion plans for the trace JIT's GPU compilers. | ||
| 4 | // | ||
| 5 | // A plan is what a GPU emitter needs to know about a traced DAG before it | ||
| 6 | // writes any code: which nodes are kernel parameters (inputs read, outputs | ||
| 7 | // written), how each parameter's buffer maps onto the trace's dominant | ||
| 8 | // (rows, cols) shape (full, a (1, cols) row, or a (1, 1) scalar), its dtype, | ||
| 9 | // and how many elements one thread moves per access. The CUDA compiler | ||
| 10 | // (trace_compiler_cuda.cpp, PTX text) and the Vulkan compiler | ||
| 11 | // (trace_compiler_vulkan.cpp, brass MIR lowered to SPIR-V) build the same two | ||
| 12 | // plans from the same DAG and differ only in how they emit them. | ||
| 13 | // | ||
| 14 | // ElementwisePlan a pure map over the dominant shape. | ||
| 15 | // RowNormPlan one row reduction (RMSNorm or LayerNorm) with the | ||
| 16 | // elementwise nodes it depends on (the "pre" chain, | ||
| 17 | // evaluated before the reduction) and the ones that depend | ||
| 18 | // on it (the "post" chain), as one block per row group. | ||
| 19 | |||
| 20 | #include "trace_dag.h" | ||
| 21 | |||
| 22 | #include <cstdint> | ||
| 23 | #include <vector> | ||
| 24 | |||
| 25 | namespace brotensor::jit::plan { | ||
| 26 | |||
| 27 | // One kernel parameter: a device pointer plus how to read through it. | ||
| 28 | struct BufferSpec { | ||
| 29 | int node_id = -1; | ||
| 30 | Dtype dtype = Dtype::FP32; | ||
| 31 | BroadcastKind bcast = BroadcastKind::Full; | ||
| 32 | }; | ||
| 33 | |||
| 34 | 8 | struct ElementwisePlan { | |
| 35 | std::vector<BufferSpec> inputs; | ||
| 36 | std::vector<BufferSpec> outputs; | ||
| 37 | 8 | int rows = 0; | |
| 38 | 8 | int cols = 0; | |
| 39 | 8 | std::int64_t numel = 0; | |
| 40 | 8 | int vec = 1; // elements per thread in the vector entry | |
| 41 | 8 | bool any_row_bcast = false; | |
| 42 | 8 | bool any_col_bcast = false; | |
| 43 | }; | ||
| 44 | |||
| 45 | // Reduction at the head of a row-norm kernel. | ||
| 46 | enum class RowReduce { RMS, Mean }; | ||
| 47 | |||
| 48 | 16 | struct RowNormPlan { | |
| 49 | ElementwisePlan ew; // buffers and shape | ||
| 50 | 8 | RowReduce reduce = RowReduce::RMS; | |
| 51 | |||
| 52 | 8 | int norm_node = -1; // the RMSNorm / LayerNorm node | |
| 53 | 8 | int x_node = -1; // what feeds it: an input, or the end of pre_ids | |
| 54 | 8 | int gamma_input = -1; // index into ew.inputs, or -1 | |
| 55 | 8 | int beta_input = -1; // index into ew.inputs, or -1 | |
| 56 | 8 | float eps = 1e-5f; | |
| 57 | |||
| 58 | // Nodes to evaluate before the reduction and after it. The pre chain runs | ||
| 59 | // in both passes — that is what lets `x += p; rms_norm(x)` fuse, since the | ||
| 60 | // reduction needs the updated x that the same kernel is about to write. | ||
| 61 | std::vector<int> pre_ids; | ||
| 62 | std::vector<int> post_ids; | ||
| 63 | |||
| 64 | // Per ew.inputs index: whether the pre chain reads it, and whether pass | ||
| 65 | // two needs it. Inputs only the post chain wants (a modulation row, say) | ||
| 66 | // stay out of pass one, and vice versa. | ||
| 67 | std::vector<char> pre_input; | ||
| 68 | std::vector<char> post_input; | ||
| 69 | |||
| 70 | // When the pre chain's result is itself one of the trace's outputs, pass | ||
| 71 | // one can store it and pass two can read it back instead of replaying the | ||
| 72 | // chain — one fewer read of every pre-chain input, and the value comes | ||
| 73 | // back out of cache. `x_output` indexes ew.outputs; -1 disables the reload. | ||
| 74 | 8 | int x_output = -1; | |
| 75 | |||
| 76 | // Per ew.outputs index: produced by the pre chain, so written in pass one. | ||
| 77 | std::vector<char> pre_output; | ||
| 78 | }; | ||
| 79 | |||
| 80 | // Bytes per element (0 for a dtype the emitters cannot address). | ||
| 81 | int elem_bytes(Dtype d); | ||
| 82 | |||
| 83 | // True when the emitters can handle this dtype at all. | ||
| 84 | bool dtype_supported(Dtype d); | ||
| 85 | |||
| 86 | // Elements per thread for a 16-byte access at this dtype. | ||
| 87 | int vec_width_for(Dtype d); | ||
| 88 | |||
| 89 | // Upper bound on a row-norm block. The actual block is rows_per_block rows of | ||
| 90 | // threads_per_row threads each, never more than this. | ||
| 91 | inline constexpr int kRowThreads = 256; | ||
| 92 | |||
| 93 | // How the row-norm kernel tiles a (rows, cols) trace at `lanes` elements per | ||
| 94 | // access. `tpr` threads cooperate on one row and `rpb` rows share a block. | ||
| 95 | // | ||
| 96 | // A row needs ceil(cols/lanes) threads to cover it in one stride. Giving it | ||
| 97 | // 256 regardless — which is what a one-row-per-block kernel does — is right | ||
| 98 | // for a 4096-wide transformer activation and catastrophic for a VAE feature | ||
| 99 | // map, where cols is 96 to 384 and rows runs into the millions: seven eighths | ||
| 100 | // of every block idles, and the launch is a million blocks deep. | ||
| 101 | // | ||
| 102 | // `rpb` is also held to a divisor of `rows`, so the "my row is past the end" | ||
| 103 | // exit is uniform across a block and the reduction's barrier is safe. | ||
| 104 | void row_norm_geometry(int cols, int rows, int lanes, int& tpr, int& rpb); | ||
| 105 | |||
| 106 | // Collects the trace's parameter buffers and rejects anything the emitters | ||
| 107 | // cannot address. Returns false rather than throwing so the caller can try a | ||
| 108 | // different fusion. | ||
| 109 | bool build_plan(const TraceDAG& dag, ElementwisePlan& plan); | ||
| 110 | |||
| 111 | // A trace is a row-norm when exactly one node is a reduction and every other | ||
| 112 | // node is elementwise. The nodes the reduction depends on become the `pre` | ||
| 113 | // chain, which the kernel evaluates in both passes; everything downstream | ||
| 114 | // becomes `post`. | ||
| 115 | bool build_row_norm_plan(const TraceDAG& dag, RowNormPlan& plan); | ||
| 116 | |||
| 117 | bool all_fp32(const TraceDAG& dag); | ||
| 118 | |||
| 119 | } // namespace brotensor::jit::plan | ||
| 120 |