src/jit/ptx_emit_elementwise.cpp
| Line | Branch | Exec | Source |
|---|---|---|---|
| 1 | #include "ptx_emit.h" | ||
| 2 | #include "ptx_emit_detail.h" | ||
| 3 | |||
| 4 | #include <sstream> | ||
| 5 | #include <string> | ||
| 6 | #include <unordered_map> | ||
| 7 | #include <vector> | ||
| 8 | |||
| 9 | namespace brotensor::jit::ptx { | ||
| 10 | |||
| 11 | using namespace brotensor::jit::ptx::detail; | ||
| 12 | |||
| 13 | namespace { | ||
| 14 | |||
| 15 | // One elementwise entry: `lanes` elements per thread, grid-stride. | ||
| 16 | ✗ | std::string emit_ew_entry(const TraceDAG& dag, const ElementwisePlan& plan, int lanes, | |
| 17 | const std::string& name) { | ||
| 18 | ✗ | Alloc a; | |
| 19 | ✗ | std::stringstream b; | |
| 20 | |||
| 21 | ✗ | const std::size_t n_out = plan.outputs.size(); | |
| 22 | ✗ | const std::size_t n_in = plan.inputs.size(); | |
| 23 | |||
| 24 | ✗ | const std::string r_n = a.R(), r_tid = a.R(), r_ntid = a.R(); | |
| 25 | ✗ | const std::string r_ctaid = a.R(), r_nctaid = a.R(); | |
| 26 | ✗ | const std::string r_i = a.R(), r_stride = a.R(); | |
| 27 | ✗ | const std::string p_oob = a.P(); | |
| 28 | |||
| 29 | ✗ | b << " ld.param.u32 " << r_n << ", [" << name << "_n];\n"; | |
| 30 | ✗ | std::vector<std::string> out_ptr(n_out), in_ptr(n_in); | |
| 31 | ✗ | for (std::size_t m = 0; m < n_out; ++m) { | |
| 32 | ✗ | out_ptr[m] = a.RD(); | |
| 33 | ✗ | b << " ld.param.u64 " << out_ptr[m] << ", [" << name << "_out_" << m << "];\n"; | |
| 34 | ✗ | } | |
| 35 | ✗ | for (std::size_t k = 0; k < n_in; ++k) { | |
| 36 | ✗ | in_ptr[k] = a.RD(); | |
| 37 | ✗ | b << " ld.param.u64 " << in_ptr[k] << ", [" << name << "_in_" << k << "];\n"; | |
| 38 | ✗ | } | |
| 39 | ✗ | b << " mov.u32 " << r_tid << ", %tid.x;\n"; | |
| 40 | ✗ | b << " mov.u32 " << r_ntid << ", %ntid.x;\n"; | |
| 41 | ✗ | b << " mov.u32 " << r_ctaid << ", %ctaid.x;\n"; | |
| 42 | ✗ | b << " mov.u32 " << r_nctaid << ", %nctaid.x;\n"; | |
| 43 | ✗ | b << " mad.lo.u32 " << r_i << ", " << r_ctaid << ", " << r_ntid << ", " << r_tid << ";\n"; | |
| 44 | ✗ | b << " mul.lo.u32 " << r_stride << ", " << r_ntid << ", " << r_nctaid << ";\n"; | |
| 45 | |||
| 46 | ✗ | b << "$L_" << name << "_head:\n"; | |
| 47 | ✗ | b << " setp.ge.u32 " << p_oob << ", " << r_i << ", " << r_n << ";\n"; | |
| 48 | ✗ | b << " @" << p_oob << " bra $L_" << name << "_exit;\n"; | |
| 49 | |||
| 50 | // Element index of this thread's first lane. | ||
| 51 | ✗ | std::string r_base = r_i; | |
| 52 | ✗ | if (lanes > 1) { | |
| 53 | ✗ | r_base = a.R(); | |
| 54 | ✗ | b << " shl.b32 " << r_base << ", " << r_i << ", " | |
| 55 | ✗ | << (lanes == 2 ? 1 : (lanes == 4 ? 2 : 3)) << ";\n"; | |
| 56 | ✗ | } | |
| 57 | |||
| 58 | // Column index, only when some operand broadcasts across rows. | ||
| 59 | ✗ | std::string r_col; | |
| 60 | ✗ | if (plan.any_row_bcast) { | |
| 61 | ✗ | r_col = a.R(); | |
| 62 | ✗ | if (is_pow2(plan.cols)) { | |
| 63 | ✗ | b << " and.b32 " << r_col << ", " << r_base << ", " << (plan.cols - 1) << ";\n"; | |
| 64 | ✗ | } else { | |
| 65 | ✗ | b << " rem.u32 " << r_col << ", " << r_base << ", " << plan.cols << ";\n"; | |
| 66 | } | ||
| 67 | ✗ | } | |
| 68 | |||
| 69 | // Row index, only when some operand broadcasts across columns. | ||
| 70 | ✗ | std::string r_row; | |
| 71 | ✗ | if (plan.any_col_bcast) { | |
| 72 | ✗ | r_row = a.R(); | |
| 73 | ✗ | if (is_pow2(plan.cols)) { | |
| 74 | ✗ | int sh = 0; | |
| 75 | ✗ | while ((1 << sh) < plan.cols) ++sh; | |
| 76 | ✗ | b << " shr.u32 " << r_row << ", " << r_base << ", " << sh << ";\n"; | |
| 77 | ✗ | } else { | |
| 78 | ✗ | b << " div.u32 " << r_row << ", " << r_base << ", " << plan.cols << ";\n"; | |
| 79 | } | ||
| 80 | ✗ | } | |
| 81 | |||
| 82 | // Byte offsets, one per (broadcast kind, element size) actually used. | ||
| 83 | ✗ | std::unordered_map<int, std::string> full_off, row_off, col_off; | |
| 84 | ✗ | auto offset_for = [&](BroadcastKind bk, Dtype dt) -> std::string { | |
| 85 | ✗ | const int esz = elem_bytes(dt); | |
| 86 | ✗ | if (bk == BroadcastKind::Scalar) return std::string(); | |
| 87 | ✗ | auto& tbl = (bk == BroadcastKind::Row) ? row_off | |
| 88 | ✗ | : (bk == BroadcastKind::Col) ? col_off : full_off; | |
| 89 | ✗ | auto it = tbl.find(esz); | |
| 90 | ✗ | if (it != tbl.end()) return it->second; | |
| 91 | ✗ | const std::string rd = a.RD(); | |
| 92 | ✗ | const std::string& idx = (bk == BroadcastKind::Row) ? r_col | |
| 93 | ✗ | : (bk == BroadcastKind::Col) ? r_row : r_base; | |
| 94 | ✗ | b << " mul.wide.u32 " << rd << ", " << idx << ", " << esz << ";\n"; | |
| 95 | ✗ | tbl[esz] = rd; | |
| 96 | ✗ | return rd; | |
| 97 | ✗ | }; | |
| 98 | |||
| 99 | ✗ | ValueMap vals; | |
| 100 | ✗ | for (std::size_t k = 0; k < n_in; ++k) { | |
| 101 | ✗ | const BufferSpec& s = plan.inputs[k]; | |
| 102 | ✗ | const bool single_lane = (s.bcast == BroadcastKind::Scalar || s.bcast == BroadcastKind::Col); | |
| 103 | ✗ | const int ln = single_lane ? 1 : lanes; | |
| 104 | ✗ | std::string addr; | |
| 105 | ✗ | if (s.bcast == BroadcastKind::Scalar) { | |
| 106 | ✗ | addr = in_ptr[k]; | |
| 107 | ✗ | } else { | |
| 108 | ✗ | const std::string off = offset_for(s.bcast, s.dtype); | |
| 109 | ✗ | addr = a.RD(); | |
| 110 | ✗ | b << " add.u64 " << addr << ", " << in_ptr[k] << ", " << off << ";\n"; | |
| 111 | ✗ | } | |
| 112 | ✗ | std::vector<std::string> v = emit_load(b, a, s.dtype, ln, addr); | |
| 113 | ✗ | if (single_lane) { | |
| 114 | // One load feeds every lane. Copy the name out first: assign() | ||
| 115 | // would clear the vector the reference points into. | ||
| 116 | ✗ | const std::string only = v[0]; | |
| 117 | ✗ | v.assign(static_cast<std::size_t>(lanes), only); | |
| 118 | ✗ | } | |
| 119 | ✗ | vals[s.node_id] = std::move(v); | |
| 120 | ✗ | } | |
| 121 | |||
| 122 | ✗ | vals = emit_chain(b, a, dag, lanes, std::move(vals)); | |
| 123 | |||
| 124 | ✗ | for (std::size_t m = 0; m < n_out; ++m) { | |
| 125 | ✗ | const BufferSpec& s = plan.outputs[m]; | |
| 126 | ✗ | const std::string off = offset_for(BroadcastKind::Full, s.dtype); | |
| 127 | ✗ | const std::string addr = a.RD(); | |
| 128 | ✗ | b << " add.u64 " << addr << ", " << out_ptr[m] << ", " << off << ";\n"; | |
| 129 | ✗ | emit_store(b, a, s.dtype, lanes, addr, vals.at(s.node_id)); | |
| 130 | ✗ | } | |
| 131 | |||
| 132 | ✗ | b << " add.u32 " << r_i << ", " << r_i << ", " << r_stride << ";\n"; | |
| 133 | ✗ | b << " bra $L_" << name << "_head;\n"; | |
| 134 | ✗ | b << "$L_" << name << "_exit:\n"; | |
| 135 | ✗ | b << " ret;\n"; | |
| 136 | |||
| 137 | ✗ | std::stringstream ss; | |
| 138 | ✗ | ss << ".visible .entry " << name << "(\n"; | |
| 139 | ✗ | for (std::size_t m = 0; m < n_out; ++m) { | |
| 140 | ✗ | ss << " .param .u64 " << name << "_out_" << m << ",\n"; | |
| 141 | ✗ | } | |
| 142 | ✗ | for (std::size_t k = 0; k < n_in; ++k) { | |
| 143 | ✗ | ss << " .param .u64 " << name << "_in_" << k << ",\n"; | |
| 144 | ✗ | } | |
| 145 | ✗ | ss << " .param .u32 " << name << "_n\n"; | |
| 146 | ✗ | ss << ") {\n"; | |
| 147 | ✗ | ss << declarations(a); | |
| 148 | ✗ | ss << b.str(); | |
| 149 | ✗ | ss << "}\n\n"; | |
| 150 | ✗ | return ss.str(); | |
| 151 | ✗ | } | |
| 152 | |||
| 153 | } // namespace | ||
| 154 | |||
| 155 | ✗ | std::string emit_elementwise(const TraceDAG& dag, const ElementwisePlan& plan, | |
| 156 | const std::string& arch) { | ||
| 157 | ✗ | std::stringstream ss; | |
| 158 | ✗ | ss << ".version 7.8\n"; | |
| 159 | ✗ | ss << ".target " << arch << "\n"; | |
| 160 | ✗ | ss << ".address_size 64\n\n"; | |
| 161 | ✗ | ss << emit_ew_entry(dag, plan, plan.vec, kEntryVec); | |
| 162 | ✗ | ss << emit_ew_entry(dag, plan, 1, kEntryScalar); | |
| 163 | ✗ | return ss.str(); | |
| 164 | ✗ | } | |
| 165 | |||
| 166 | } // namespace brotensor::jit::ptx | ||
| 167 |