Skip to content

[Ascend] Port pto-dev PTO/VMI increments onto the PR #3310 baseline - #1

Open
jimmychou0 wants to merge 19 commits into
mainfrom
pto-dev
Open

jimmychou0 wants to merge 19 commits into
mainfrom
pto-dev

Conversation

@jimmychou0

@jimmychou0 jimmychou0 commented Sep 30, 2026 •

Copy link
Copy Markdown
Collaborator

Summary

Ports the pto-dev PTO/VMI increments from WenboCodes/tilelang-deepseek pto-dev onto the merged PR tile-ai#3310 baseline: 47 commits from fa11da8f ([PTO] Add PTO codegen backend support tile-ai#140) through 34c6d3ea, plus 7 port-integration fixes (latest 3d70ede9).

What's included

  • PTO codegen: SIMT and CUBE kernels, GEMM BF16/FP32/HF32/FP8, blockscaled/MX, multi-kernel JIT modules
  • VMI language surface: T.vmi, vzip/vunzip, vgather, packed FP4 vcvt, width-changing vinterpret_cast
  • SIMT allreduce, atomics with correct signedness, pto.vecscope, unroll-hint forwarding, PTO perf-regression CI job

Port-integration fixes (validated on hardware)

  1. 50bb9af9 relax ascend_load_cbuf_to_ca/cb to variable arg count (9/16 MX companion)
  2. e6ae351c point PTO codegen at the relocated codegen_py sources
  3. d4d04fb7 pass compile_format to BinaryCache.load/save
  4. 90054e32 restore 7 pre-fa11da8f PTO example modules
  5. 0774ecff restore T.blockscaled_gemm (delegates sfa/sfb to the 16-slot tl.tileop.gemm_blockscaled protocol)
  6. 3d70ede9 restore the scale= MX companion kwarg on T.copy

Validation (238, x86_64 + Ascend950PR; pto-as vmi-v0.1.9 b9da12032; CANN 9.1 runtime + 9.2 OPP; py3.11/torch 2.10 — deviation from the aarch64 CI env noted)

job1 (-m "not pto" ../examples/ascend)    88 passed, 21 failed*
job2 (-m "not pto" ascend)              1181 passed,  0 failed  ✓
job3 (-m "pto"     ../examples/ascend)    73 passed, 17 failed*
job4 (-m "pto"     ascend)               219 passed,  0 failed  ✓

*Known issue A — pto-as mad_raw lowering regression (31 cases)

All test_blockscaled_gemm_l0_*, test_gemm_padding_irregular_shapes (fp8 blockscaled) and test_gemm_splitk/test_gemm_ub_merge failures are pto.mad → pto.mad_raw "explicitly marked illegal" in pto-as v0.1.8..master (tested e61df455d, stack-v3 98a338585, stack7, b9da12032). The identical generated .pto lowers fine on the pre-0915 PTOAS build, and the .pto text is byte-identical to what original pto-dev emits — no tilelang-side regression. Needs a pto-as fix or a matched release; tracked separately.

*Known issue B — cross-file state contamination (~7 cases, non-deterministic)

test_vecadd[simd/simt], test_while_pipelined, test_int64_stride_vectorize_load, test_sync_inter_autoschedule, test_per_token_cast_to_fp8 fail only when mixed with other files in one pytest process; each passes standalone and per-file (test_vecadd.py whole file: 2 passed). Test-order-sensitive runtime state; recommend per-file isolation in CI while investigating.

Notes

jimmychou0 and others added 17 commits October 26, 2026 10:39
Add the VMI (virtual machine instruction) intrinsics and PTO codegen:

- tilelang/ascend/language/vmi.py: the T.vmi.* surface, exposed as T.vmi
  through tilelang.ascend.language - vector loads/stores, masks,
  arithmetic/bitwise/shift ops, vcvt (incl. packed FP4 sources),
  vinterpret_cast, vzip/vunzip, precision-aware vdiv, gather/scatter,
  histogram ops, and register allocation helpers
- 54 tl.vmi.* TIR builtins registered in src/op/builtin.{h,cc}
- PTO codegen: pto.vmi.* emission with typed-literal wrapping, FP4
  pair expansion, eager vreg allocation for vector local.var, the
  signed/unsigned integer type split (pto.si*/pto.ui*), and a clear
  AscendC rejection diagnostic; generated kernels import
  tilelang.contrib.ptodsl before ptodsl so this build's native
  libraries load ahead of the PTODSL MLIR modules
- KernelParam.storage_shape: packed-storage shapes for sub-byte dtypes

Co-authored-by: KurrinQu <qukelin1991@163.com>
Co-authored-by: Jay Zhuang <80731350+learning-chip@users.noreply.github.com>
Co-authored-by: Zhang Zhendong <zhendong404@gmail.com>
Co-authored-by: cj <erhsh_165@126.com>
Add the T.Parallel -> PTO VMI vectorization pipeline:

- Verify -> Legalize -> Vectorize pass pipeline ahead of
  AscendSimdVFLowerParallel, with pass registration, the Python
  interface, and routing tests (SIMT and hand-written VMI regions are
  bypassed inside the passes; control-flow rewriting stays out of scope)
- Verify: shared region scan and inverse lane-mapping analysis, the
  LayoutInference-for-PTO branch, 1D/2D lane mapping, and
  T.SimdVF(lanes=...) frontend support, with contract tests
- Vectorize: chunk loops with dynamic prefix tail masks, remaining
  state reset per row/column and per sequential unit, and dtype/cast
  codegen for element-wise add/sub/mul/div, unary abs/exp/log/sqrt,
  and f32 <-> f16/bf16 conversions; runtime lane-uniform float
  computations are rejected (they would stay scalar inside SIMD_VF)
- Device tests: element-wise, tail-mask, and FP8 suites

Co-authored-by: qukelin <qukelin1991@163.com>
Co-authored-by: jimmychou <47636600+jimmychou0@users.noreply.github.com>
The si32 reducer strategy serves both signed si32 and signless i32
carriers, but created its initial-value constant as signless pto.i32.
With TIR signedness now carried into PTODSL, scalar.select mixing an
si32 carrier with an i32 constant is rejected by current PTOAS
(pto.select integer arms require matching types).

Both creation points (_emit_warp_hw_reduce and _emit_cross_warp_reduce)
now share _identity_const: the strategy key still selects the constant's
numeric value, while its integer type follows the carrier value's own
signedness (signed carriers get pto.si32, signless keep pto.i32).
Signed-off-by: wangchen <wangchen965@huawei.com>
wangchen and others added 2 commits October 27, 2026 09:12
Signed-off-by: wangchen <wangchen965@huawei.com>
feat(pto): default disable LegalizePTOVMIMemoryAccess
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants