feat(ptodsl): Add a5 templates for tconcat, tdequant, tprint, treshape#968
feat(ptodsl): Add a5 templates for tconcat, tdequant, tprint, treshape#968Crystal-wzy wants to merge 0 commit into
Conversation
There was a problem hiding this comment.
Code Review
This pull request introduces several new tile operations to the ptodsl framework, including concat, dequant, print, and reshape, along with their corresponding A5 hardware templates, documentation, and system tests. Feedback on the implementation highlights an optimization opportunity in the tdequant templates (template_tdequant_i16 and template_tdequant_i8), where hoisting the loop-invariant load and broadcast of the per-row scale and offset coefficients outside of the column loops would avoid redundant operations and improve vector pipeline efficiency.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
| with pto.for_(0, valid_rows, step=1) as row: | ||
| remained = valid_cols | ||
| col_loop = pto.for_(0, valid_cols, step=lanes).carry(remained=remained) | ||
| with col_loop: | ||
| col = col_loop.iv | ||
| mask, remained = pto.make_mask(pto.f32, col_loop.remained) | ||
| # i16 -> f32 in one shot (unpack-on-load + even-part convert). | ||
| src_v = pto.vlds(src[row, col:], dist="UNPK_B16") | ||
| value = pto.vcvt(src_v, pto.f32, src_full, part=pto.VcvtPartMode.EVEN) | ||
| offset_b = _broadcast_row(offset, row, mask) | ||
| scaled = pto.vmul(pto.vsub(value, offset_b, mask), _broadcast_row(scale, row, mask), mask) | ||
| pto.vsts(scaled, dst[row, col:], mask) | ||
| col_loop.update(remained=remained) |
There was a problem hiding this comment.
The load and broadcast of the per-row scale and offset coefficients are loop-invariant with respect to the column loop (col_loop). Hoisting these operations outside of col_loop avoids redundant vlds and vdup instructions in every iteration, significantly improving vector pipeline efficiency.
mask_f32_full = pto.make_mask(pto.f32, pto.PAT.ALL)\n with pto.for_(0, valid_rows, step=1) as row:\n scale_b = _broadcast_row(scale, row, mask_f32_full)\n offset_b = _broadcast_row(offset, row, mask_f32_full)\n remained = valid_cols\n col_loop = pto.for_(0, valid_cols, step=lanes).carry(remained=remained)\n with col_loop:\n col = col_loop.iv\n mask, remained = pto.make_mask(pto.f32, col_loop.remained)\n # i16 -> f32 in one shot (unpack-on-load + even-part convert).\n src_v = pto.vlds(src[row, col:], dist=\"UNPK_B16\")\n value = pto.vcvt(src_v, pto.f32, src_full, part=pto.VcvtPartMode.EVEN)\n scaled = pto.vmul(pto.vsub(value, offset_b, mask), scale_b, mask)\n pto.vsts(scaled, dst[row, col:], mask)\n col_loop.update(remained=remained)| with pto.for_(0, valid_rows, step=1) as row: | ||
| remained = valid_cols | ||
| next_remained = valid_cols - lanes_i32 | ||
| col_loop = pto.for_(0, valid_cols, step=lanes_i16).carry( | ||
| remained=remained, next_remained=next_remained, | ||
| ) | ||
| with col_loop: | ||
| col = col_loop.iv | ||
| # i8 -> i32 sign-extending interleave (mirrors template_tcvt_si8_to_i32). | ||
| mask_b16_cur, remained = pto.make_mask(pto.i16, col_loop.remained) | ||
| mask_b16_next, next_remained = pto.make_mask(pto.i16, col_loop.next_remained) | ||
| mask_b32_cur = pto.punpack(mask_b16_cur, pto.PredicatePart.LOWER, to_type=pto.mask_b32) | ||
| mask_b32_next = pto.punpack(mask_b16_next, pto.PredicatePart.LOWER, to_type=pto.mask_b32) | ||
| vec_si8_0 = pto.vlds(src[row, col:], dist="UNPK_B8") | ||
| vec_ui8_1, vec_ui8_2 = pto.vintlv(pto.vbitcast(vec_si8_0, pto.ui8), v_zero) | ||
| i32_cur = pto.vcvt(pto.vbitcast(vec_ui8_1, pto.si8), pto.i32, b8_mask, part=pto.VcvtPartMode.P0) | ||
| i32_next = pto.vcvt(pto.vbitcast(vec_ui8_2, pto.si8), pto.i32, b8_mask, part=pto.VcvtPartMode.P0) | ||
| # i32 -> f32 (rnd=Z, as in C++). | ||
| value_cur = pto.vcvt(i32_cur, pto.f32, mask_b32_cur, rnd=pto.VcvtRoundMode.Z) | ||
| value_next = pto.vcvt(i32_next, pto.f32, mask_b32_next, rnd=pto.VcvtRoundMode.Z) | ||
| scaled_cur = pto.vmul( | ||
| pto.vsub(value_cur, _broadcast_row(offset, row, mask_b32_cur), mask_b32_cur), | ||
| _broadcast_row(scale, row, mask_b32_cur), | ||
| mask_b32_cur, | ||
| ) | ||
| scaled_next = pto.vmul( | ||
| pto.vsub(value_next, _broadcast_row(offset, row, mask_b32_next), mask_b32_next), | ||
| _broadcast_row(scale, row, mask_b32_next), | ||
| mask_b32_next, | ||
| ) | ||
| pto.vsts(scaled_cur, dst[row, col:], mask_b32_cur) | ||
| pto.vsts(scaled_next, dst[row, col + lanes_i32:], mask_b32_next) | ||
| col_loop.update(remained=remained, next_remained=next_remained) |
There was a problem hiding this comment.
Similar to the i16 template, the per-row scale and offset coefficients are loop-invariant with respect to col_loop. Hoisting their load and broadcast outside of the column loop avoids four redundant vlds and vdup operations per iteration, which greatly reduces instruction overhead.
mask_f32_full = pto.make_mask(pto.f32, pto.PAT.ALL)\n with pto.for_(0, valid_rows, step=1) as row:\n scale_b = _broadcast_row(scale, row, mask_f32_full)\n offset_b = _broadcast_row(offset, row, mask_f32_full)\n remained = valid_cols\n next_remained = valid_cols - lanes_i32\n col_loop = pto.for_(0, valid_cols, step=lanes_i16).carry(\n remained=remained, next_remained=next_remained,\n )\n with col_loop:\n col = col_loop.iv\n # i8 -> i32 sign-extending interleave (mirrors template_tcvt_si8_to_i32).\n mask_b16_cur, remained = pto.make_mask(pto.i16, col_loop.remained)\n mask_b16_next, next_remained = pto.make_mask(pto.i16, col_loop.next_remained)\n mask_b32_cur = pto.punpack(mask_b16_cur, pto.PredicatePart.LOWER, to_type=pto.mask_b32)\n mask_b32_next = pto.punpack(mask_b16_next, pto.PredicatePart.LOWER, to_type=pto.mask_b32)\n vec_si8_0 = pto.vlds(src[row, col:], dist=\"UNPK_B8\")\n vec_ui8_1, vec_ui8_2 = pto.vintlv(pto.vbitcast(vec_si8_0, pto.ui8), v_zero)\n i32_cur = pto.vcvt(pto.vbitcast(vec_ui8_1, pto.si8), pto.i32, b8_mask, part=pto.VcvtPartMode.P0)\n i32_next = pto.vcvt(pto.vbitcast(vec_ui8_2, pto.si8), pto.i32, b8_mask, part=pto.VcvtPartMode.P0)\n # i32 -> f32 (rnd=Z, as in C++).\n value_cur = pto.vcvt(i32_cur, pto.f32, mask_b32_cur, rnd=pto.VcvtRoundMode.Z)\n value_next = pto.vcvt(i32_next, pto.f32, mask_b32_next, rnd=pto.VcvtRoundMode.Z)\n scaled_cur = pto.vmul(\n pto.vsub(value_cur, offset_b, mask_b32_cur),\n scale_b,\n mask_b32_cur,\n )\n scaled_next = pto.vmul(\n pto.vsub(value_next, offset_b, mask_b32_next),\n scale_b,\n mask_b32_next,\n )\n pto.vsts(scaled_cur, dst[row, col:], mask_b32_cur)\n pto.vsts(scaled_next, dst[row, col + lanes_i32:], mask_b32_next)\n col_loop.update(remained=remained, next_remained=next_remained)675f215 to
e8a417d
Compare
Add four A5 TileLib templates plus their op bindings, catalog registration, tests, and docs so the PTODSL tile library can lower graphs containing these pto ops on A5.
Summary
candidatesattr) and ExpandTileOp (reads it), and drops the attr on non-fusible ops like tconcat, so ExpandTileOp finds no candidate. ST kernels are single-op (nothing to fuse), so native_build.py opts out via PTODSL_DISABLE_OP_FUSION (set in the ci_sim.yml TileLib ST step), passing --enable-op-fusion=false to ptoas.Testing