Spacemit backend bugfixes and support tle.raw - #1
Open
zuoweixia497 wants to merge 16 commits into
Open
zuoweixia497 wants to merge 16 commits into
zuoweixia497 wants to merge 16 commits into
Conversation
The spec overlay copy of code_generator.py has drifted from the main python/triton/compiler/code_generator.py (still references the old language.extra.smt structure) and is no longer maintained. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
Python side (language/spine_raw/): AST-visitor codegen producing linalg/vector MLIR from a restricted Python subset, the call_registry that records spine_raw.call() during JIT tracing, and the LLVM-direct text codegen path. C++ side: new tle.dsl_region op (TLEOps.td) whose region holds the parsed raw kernel body; DSLRegionOpPattern (TLEToLinalg) parses the raw_linalg text, traces operands through the ptr.to_ptr / memref.reinterpret_cast / memref.cast chain back to the original memref block args, and creates spine_ext.raw_region. The ptr pipeline (TritonToUnstructured) skips DSLRegionOp so its !tt.ptr operands stay untouched. triton_shared.cc gains the create_tle_dsl_region builder binding that parses the raw MLIR text at build time. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
- Pass --allow-unregistered-dialect to spine-triton-opt so the conversion can create spine_ext.raw_region in generic form. - Rewrite bit-packed dense i1 attribute blobs to the byte-per-element form spine-opt parses, fixing the cross-tool linalg IR handoff. - LLVM-direct bypass: translate a pending llvm.func module straight to LLVM IR via mlir-translate, skipping spine-opt. - Mixed-mode: graft llvm.func siblings and host-to-sibling i64 call bridges into the lowered ll.mlir (post spine-opt, where the host is uniform llvm dialect and memrefs are already descriptors). Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
…th dialect ops The 15 cpu libdevice shims (acos/atan/tan/asin/cosh/exp2/expm1/log2/ log10/log1p/sinh/acosh/asinh/atanh/cbrt) called non-existent _semantic.create_X helpers; reimplement them with core.extern + extern_elementwise. ffs gains CUDA semantics (1-based, 0 for 0) via math.cttz + select, and pow promotes mixed dtypes before dispatch. ConversionPatterns.hpp: whitelist those symbols in ConvertExternSpecialMath and dispatch them to the matching math.* ops, plus linalg.rint -> math.roundeven (rint is round-half-to-even, distinct from math.round). Without the mapping the tt.extern_elementwise ops leak into linalg and spine-opt rejects them as an unregistered dialect. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
Pointer addresses containing subtraction (e.g. ih = oh*s - pad + kh*d in conv im2col patterns) hit unsupported arith.subi and fell off the structured path. Handle SubIOp as AddI(a, -b) in both the structured and unstructured analyses, subtracting offsets/strides via subOFRs. subOFRs now unifies operand types before creating arith.subi: the unstructured path stores raw tensor values in offsets (rebuildAsUnsupportedOp), and index-typed values mixed with tensor operands previously produced an invalid arith.subi. Also document why UseAnalysis no longer overrides visitNonControlFlowArguments on LLVM 22. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
…ulation Non-gather dimensions must also contribute to the flat target offset: PtrAnalysis has already folded the memory stride into each dimension's offset scalar (e.g. pid_m * N), so skipping them dropped the base offset and gathered from the wrong rows. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
Scalar i1 masks (e.g. tl.store(ptr, v, mask=pid==0)) were dropped twice on the way down: (a) PtrAnalysis rewriteLoad/StoreOp parsed the predicate into an empty-dims MaskState and emitted unconditional affine ops -- now bails out for non-shaped masks so the unstructured path keeps the predicate; (b) the scalar Gather/Scatter converters in UnstructuredToMemref ignored tts.gather/scatter's mask operand -- now guarded with scf.if, yielding other/zero on the load's false branch. GetStructuredStateOp rewrite failure only remapped result #0, leaving the decomposed offset/stride results live so the op leaked into the final linalg and spine-opt rejected it as an unregistered dialect. Revert fully instead (result #0 -> original value, the rest -> index zero constants) and erase; the pass-side walk no longer touches the op afterwards (use-after-free). MaskAnalysis: parseConstant now carries a shaped constant's shape in dims like tt.splat (empty dims caused out-of-bounds SmallVector access in parseCmp/minStateScalar/parseBroadcast/parseExpandDims), and parseExpandDims/parseBroadcast bail out defensively on rank mismatch. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
New loop-ptr-carry-to-offset pass: rewrites scf.while loops carrying scalar !tt.ptr values advanced by loop-invariant strides (ptr += stride) into integer-offset-carrying loops, and flattens pointer merging if/else diamonds into select chains (IfPtrYieldToSelectPattern) -- loop-carried raw pointers and cf.br block args of !ptr.ptr type are not legalized downstream. Registered as a standalone pass and wired into the experimental pipeline before pointer analysis. Pipeline fixes in TritonToLinalgExperimentalPass: - erase gpu.barrier (tl.debug_barrier): CTA-scope sync is meaningless when each CPU program executes sequentially, and bufferization rejects its unknown memory side effects; - skip remove-dead-values: LLVM 22 corrupts scf.for loops whose iter-args are live in the body but dead outside (K-loop matmuls fail verification with mismatched inits/iter_args); - inline spine_ext.raw_region bodies produced by TLEToLinalg so spine-opt's e2e pipeline receives clean linalg/memref/vector IR, followed by a second XSMTToLinalg round for the inlined proton ops. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
Multiple tensor args sharing one host storage (e.g. the two output views of torch.polar) each got an independent remote buffer; reading them back in order made the last readback clobber the kernel's writes through the earlier arg's buffer. Cache remote addresses by host storage data_ptr so each storage is uploaded and read back exactly once. Upload now copies the raw buffer via ctypes.string_at instead of iterating bytes(storage) element-by-element (~seconds for 4 MiB). TypedPtr args without an allocation size derive it from the following scalar element count. clear_cache memsets locally: routing zero_() through the Triton override uploads and reads back the full benchmark cache buffer through RPC on every repetition. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
Hardcoded size=1 reinterpretations of unranked memrefs truncated tensors (masked_select's tensor<256xi8> was seen as memref<1xi8>). Use dynamic size with a large access bound instead, and keep the block-argument path in PtrAnalysis::getScalarMemRef consistent. NOTE: work-in-progress, not end-to-end verified -- the masked_select failures it targets are currently blocked earlier (downstream ConvertSpeStructToVector assertion on integer generics), so this change's effect is unobservable until that is fixed. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
Copy spine-triton python/tests/raw into the spacemit backend examples so the raw eDSL test suite travels with the plugin.
This reverts commit 597cdbc.
Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
Keep only test_raw_mv_svector.py and test_raw_mv_three_stage.py, which pass end-to-end (42/42 and 16/16) in the CI environment (venv + QEMU RPC). Remove the other 41 exploratory benches, diagnostics and test files. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
The FlagTree plugin installs the backend as triton.backends.spacemit; the kept raw mv tests were still importing from the spine_triton tree, which only works with a sys.modules alias. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
alex-spacemit
pushed a commit
that referenced
this pull request
Sep 22, 2026
…integer bitcast (flagos-ai#1240) * [amd] CanonicalizePointers: merge scf.if fat-ptr attrs instead of asserting ConvertSCFIfOp required the then/else branches of an scf.if to carry fat pointers with identical attributes (canNarrow), asserting equality and propagating only the then-branch attrs. A conditional pointer select whose arms differ (e.g. FlagGems cat/div_tensor kernels) tripped the assertion and crashed the compiler with "PassManager::run failed" on assert-enabled builds. Port the upstream Triton behavior: add FatPtrAttrs::intersect, drop the strict equality assert, and merge the then/else attributes conservatively (canNarrow only when both arms allow it). Fixes cat, concatenate and div_tensor on AMD. * [amd] AtomicCAS: bitcast non-integer element types around cmpxchg llvm.cmpxchg only accepts integer or pointer operands, but the AMD AtomicCAS lowering fed f16/bf16 values straight to LLVM::AtomicCmpXchgOp, so f16/bf16 tl.atomic_cas (e.g. FlagGems scatter_reduce with reduce="multiply") failed to lower: "'llvm.cmpxchg' op operand #1 must be signless integer or LLVM pointer type, but got 'f16'". Bitcast the cmp/val to an integer of the same bit width before the atomic and bitcast the loaded result back, mirroring upstream. Also extract the loaded value from field 0 of the {value, success} pair in the tensor path (was indexed by the loop counter). Fixes scatter_reduce on AMD. * style: clang-format AtomicCAS extract_val * [amd] docs: trace CanonicalizePointers/AtomicCAS fixes to upstream Triton Add provenance comments to the two AMD backend fixes so reviewers can follow them back to the corresponding upstream Triton work: - CanonicalizePointers scf.if fat-ptr merge: mirrors upstream handling of merge points with different pointer promotability (triton-lang/triton PR #9891, issue #9859); the 3.6.0 base instead asserted equality. - AtomicCAS non-integer bitcast + field-0 extract: brings the AMD backend in line with upstream PR #9116 (bitcast non-int cmpxchg operands) and PR #8867 (correct {value,success} struct index), both of which the 3.6.0 base predates. Comment-only; no functional change.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
[SpacemiT] Add spine_raw eDSL, backend lowering/IR fixes
概述
本 PR 集中了 SpacemiT 后端（
third_party/spacemit）的两方面工作：@triton.jitkernel 内以受限 Python 子集编写 raw kernel；1. spine_raw eDSL
动机：Triton 结构化路径无法表达 RVV/vector_ext 的底层细节（如
llvm.riscv.vfwmacc矩阵指令、scalable vector 布局）。spine_raw eDSL 允许在@triton.jithost kernel 内以受限 Python 子集直接编写 raw kernel，经 AST visitor 生成 linalg/vector MLIR，通过tle.dsl_region注入 TTIR metadata，再由 spine-opt 的 InlinePass 展开进主流水线，全程与 Triton 前端共存、无模板机制。示例（
python/examples/raw/）：保留两个端到端验证过的测试——test_raw_mv_svector.py（纯 svector 与 svector+pack 两种写法，42 例）、test_raw_mv_three_stage.py（三阶段融合 mv：svector pre/post-scale +llvm.riscv.vfwmacccall_intrinsic，多核 grid 切分，16 例）。2. 后端 lowering / IR 修复
针对最新 FlagGems 算子库进行了 bugfix 修复支持。修复后在 QEMU RPC（libspert 0.6.3）环境下对 FlagGems 全量算子测试做了验证：368 个测试文件中 277 个通过（约 75%），覆盖 mm/addmm/bmm、conv2d/conv3d、softmax、layer_norm、cat/concatenate/stack 等主流算子；