【求助】使用 spine-triton 的 smt.parallal 报错

使用 spine-triton 的 smt.parallal 时会报编译错误:

spine-opt: /home/share/nfs_share/llvm-proj/llvm/mlir/include/mlir/IR/UseDefLists.h:198: mlir::IRObjectWithUseList<OperandType>::~IRObjectWithUseList() [with OperandType = mlir::OpOperand]: Assertion `use_empty() && "Cannot destroy a value that still has uses!"' failed.
PLEASE submit a bug report to https://github.com/llvm/llvm-project/issues/ and include the crash backtrace and instructions to reproduce the bug.
Stack dump:
0.      Program arguments: /home/zhaoyy/triton-exp/.venv/lib/python3.13/site-packages/triton/backends/spine_triton/bin/spine-opt /tmp/tmpob36d_xd/linalg.mlir "--spine-triton-e2e-pipeline=enable-always-tls=0 enable-fuse-group=false" -o /tmp/tmpob36d_xd/ll.mlir
 #0 0x0000000004f241d0 llvm::sys::PrintStackTrace(llvm::raw_ostream&, int) (/home/zhaoyy/triton-exp/.venv/lib/python3.13/site-packages/triton/backends/spine_triton/bin/spine-opt+0x4f241d0)
 #1 0x0000000004f21da8 SignalHandler(int, siginfo_t*, void*) Signals.cpp:0:0
 #2 0x0000003f84770620 (linux-vdso.so.1+0x620)
 #3 0x0000003f842f14aa (/usr/lib/riscv64-linux-gnu/libc.so.6+0x794aa)
 #4 0x0000003f842b2320 raise (/usr/lib/riscv64-linux-gnu/libc.so.6+0x3a320)
 #5 0x0000003f8429fa76 abort (/usr/lib/riscv64-linux-gnu/libc.so.6+0x27a76)
 #6 0x0000003f842e760c (/usr/lib/riscv64-linux-gnu/libc.so.6+0x6f60c)
 #7 0x0000003f842abd4c __assert_fail (/usr/lib/riscv64-linux-gnu/libc.so.6+0x33d4c)
 #8 0x0000000004ba7394 mlir::Block::eraseArguments(unsigned int, unsigned int) (/home/zhaoyy/triton-exp/.venv/lib/python3.13/site-packages/triton/backends/spine_triton/bin/spine-opt+0x4ba7394)
 #9 0x0000000000a94b82 (anonymous namespace)::ParallelRewriter::matchAndRewrite(mlir::speir::spert::ParallelOp, mlir::speir::spert::ParallelOpAdaptor, mlir::ConversionPatternRewriter&) const SpeStreamtoLLVM.cc:0:0
#10 0x0000000000a95ad8 llvm::LogicalResult mlir::ConversionPattern::dispatchTo1To1<mlir::ConvertOpToLLVMPattern<mlir::speir::spert::ParallelOp, false>, mlir::speir::spert::ParallelOp>(mlir::ConvertOpToLLVMPattern<mlir::speir::spert::ParallelOp, false> const&, mlir::speir::spert::ParallelOp, mlir::speir::spert::ParallelOp::GenericAdaptor<llvm::ArrayRef<mlir::ValueRange>>, mlir::ConversionPatternRewriter&) (/home/zhaoyy/triton-exp/.venv/lib/python3.13/site-packages/triton/backends/spine_triton/bin/spine-opt+0xa95ad8)
#11 0x0000000000a95b8c mlir::ConvertOpToLLVMPattern<mlir::speir::spert::ParallelOp, false>::matchAndRewrite(mlir::Operation*, llvm::ArrayRef<mlir::ValueRange>, mlir::ConversionPatternRewriter&) const (/home/zhaoyy/triton-exp/.venv/lib/python3.13/site-packages/triton/backends/spine_triton/bin/spine-opt+0xa95b8c)
#12 0x00000000049a2a98 mlir::ConversionPattern::matchAndRewrite(mlir::Operation*, mlir::PatternRewriter&) const (/home/zhaoyy/triton-exp/.venv/lib/python3.13/site-packages/triton/backends/spine_triton/bin/spine-opt+0x49a2a98)
#13 0x00000000049e86fe mlir::PatternApplicator::matchAndRewrite(mlir::Operation*, mlir::PatternRewriter&, llvm::function_ref<bool (mlir::Pattern const&)>, llvm::function_ref<void (mlir::Pattern const&)>, llvm::function_ref<llvm::LogicalResult (mlir::Pattern const&)>)::'lambda'()::operator()() const PatternApplicator.cpp:0:0
#14 0x00000000049e951c mlir::PatternApplicator::matchAndRewrite(mlir::Operation*, mlir::PatternRewriter&, llvm::function_ref<bool (mlir::Pattern const&)>, llvm::function_ref<void (mlir::Pattern const&)>, llvm::function_ref<llvm::LogicalResult (mlir::Pattern const&)>) (/home/zhaoyy/triton-exp/.venv/lib/python3.13/site-packages/triton/backends/spine_triton/bin/spine-opt+0x49e951c)
#15 0x000000000499f080 (anonymous namespace)::OperationLegalizer::legalize(mlir::Operation*) DialectConversion.cpp:0:0
#16 0x000000000499ff06 mlir::OperationConverter::convert(mlir::Operation*, bool) (/home/zhaoyy/triton-exp/.venv/lib/python3.13/site-packages/triton/backends/spine_triton/bin/spine-opt+0x499ff06)
#17 0x00000000049a466a mlir::OperationConverter::applyConversion(llvm::ArrayRef<mlir::Operation*>) (/home/zhaoyy/triton-exp/.venv/lib/python3.13/site-packages/triton/backends/spine_triton/bin/spine-opt+0x49a466a)
#18 0x00000000049a6402 applyConversion(llvm::ArrayRef<mlir::Operation*>, mlir::ConversionTarget const&, mlir::FrozenRewritePatternSet const&, mlir::ConversionConfig, (anonymous namespace)::OpConversionMode) DialectConversion.cpp:0:0
#19 0x00000000049a6516 mlir::applyPartialConversion(mlir::Operation*, mlir::ConversionTarget const&, mlir::FrozenRewritePatternSet const&, mlir::ConversionConfig) (/home/zhaoyy/triton-exp/.venv/lib/python3.13/site-packages/triton/backends/spine_triton/bin/spine-opt+0x49a6516)
#20 0x0000000000a8ed66 (anonymous namespace)::ConvertSpeStreamToLLVM::runOnOperation() SpeStreamtoLLVM.cc:0:0
#21 0x0000000004a9d1b4 mlir::detail::OpToOpPassAdaptor::run(mlir::Pass*, mlir::Operation*, mlir::AnalysisManager, bool, unsigned int) (/home/zhaoyy/triton-exp/.venv/lib/python3.13/site-packages/triton/backends/spine_triton/bin/spine-opt+0x4a9d1b4)
#22 0x0000000004a9d4d2 mlir::detail::OpToOpPassAdaptor::runPipeline(mlir::OpPassManager&, mlir::Operation*, mlir::AnalysisManager, bool, unsigned int, mlir::PassInstrumentor*, mlir::PassInstrumentation::PipelineParentInfo const*) (/home/zhaoyy/triton-exp/.venv/lib/python3.13/site-packages/triton/backends/spine_triton/bin/spine-opt+0x4a9d4d2)
#23 0x0000000004a9e69c mlir::PassManager::run(mlir::Operation*) (/home/zhaoyy/triton-exp/.venv/lib/python3.13/site-packages/triton/backends/spine_triton/bin/spine-opt+0x4a9e69c)
#24 0x0000000000a1f688 performActions(llvm::raw_ostream&, std::shared_ptr<llvm::SourceMgr> const&, mlir::MLIRContext*, mlir::MlirOptMainConfig const&) MlirOptMain.cpp:0:0
#25 0x0000000000a2014c processBuffer(llvm::raw_ostream&, std::unique_ptr<llvm::MemoryBuffer, std::default_delete<llvm::MemoryBuffer>>, llvm::MemoryBufferRef, mlir::MlirOptMainConfig const&, mlir::DialectRegistry&, mlir::SourceMgrDiagnosticVerifierHandler*, llvm::ThreadPoolInterface*) MlirOptMain.cpp:0:0
#26 0x0000000000a2036a llvm::LogicalResult llvm::function_ref<llvm::LogicalResult (std::unique_ptr<llvm::MemoryBuffer, std::default_delete<llvm::MemoryBuffer>>, llvm::MemoryBufferRef const&, llvm::raw_ostream&)>::callback_fn<mlir::MlirOptMain(llvm::raw_ostream&, std::unique_ptr<llvm::MemoryBuffer, std::default_delete<llvm::MemoryBuffer>>, mlir::DialectRegistry&, mlir::MlirOptMainConfig const&)::'lambda'(std::unique_ptr<llvm::MemoryBuffer, std::default_delete<llvm::MemoryBuffer>>, llvm::MemoryBufferRef, llvm::raw_ostream&)>(long, std::unique_ptr<llvm::MemoryBuffer, std::default_delete<llvm::MemoryBuffer>>, llvm::MemoryBufferRef const&, llvm::raw_ostream&) MlirOptMain.cpp:0:0
#27 0x0000000004c67ec8 mlir::splitAndProcessBuffer(std::unique_ptr<llvm::MemoryBuffer, std::default_delete<llvm::MemoryBuffer>>, llvm::function_ref<llvm::LogicalResult (std::unique_ptr<llvm::MemoryBuffer, std::default_delete<llvm::MemoryBuffer>>, llvm::MemoryBufferRef const&, llvm::raw_ostream&)>, llvm::raw_ostream&, llvm::StringRef, llvm::StringRef) (/home/zhaoyy/triton-exp/.venv/lib/python3.13/site-packages/triton/backends/spine_triton/bin/spine-opt+0x4c67ec8)
#28 0x0000000000a18ca6 mlir::MlirOptMain(llvm::raw_ostream&, std::unique_ptr<llvm::MemoryBuffer, std::default_delete<llvm::MemoryBuffer>>, mlir::DialectRegistry&, mlir::MlirOptMainConfig const&) (.part.0) MlirOptMain.cpp:0:0
#29 0x0000000000a206c0 mlir::MlirOptMain(int, char**, llvm::StringRef, llvm::StringRef, mlir::DialectRegistry&) (/home/zhaoyy/triton-exp/.venv/lib/python3.13/site-packages/triton/backends/spine_triton/bin/spine-opt+0xa206c0)
#30 0x0000000000a20836 mlir::MlirOptMain(int, char**, llvm::StringRef, mlir::DialectRegistry&) (/home/zhaoyy/triton-exp/.venv/lib/python3.13/site-packages/triton/backends/spine_triton/bin/spine-opt+0xa20836)
#31 0x00000000008cbc60 main (/home/zhaoyy/triton-exp/.venv/lib/python3.13/site-packages/triton/backends/spine_triton/bin/spine-opt+0x8cbc60)
#32 0x0000003f8429ff62 (/usr/lib/riscv64-linux-gnu/libc.so.6+0x27f62)
#33 0x0000003f842a004e __libc_start_main (/usr/lib/riscv64-linux-gnu/libc.so.6+0x2804e)
#34 0x0000000000995cfc _start /builds/compiler/riscv-gnu-toolchain/glibc/csu/../sysdeps/riscv/start.S:68:0
Aborted
Traceback (most recent call last):
  File "/home/zhaoyy/triton-exp/test_smt_mm.py", line 248, in <module>
    _run_correctness(M, N, K, dtype=dtype, device=device)
    ~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/zhaoyy/triton-exp/test_smt_mm.py", line 195, in _run_correctness
    C = triton_mm(A, B)
  File "/home/zhaoyy/triton-exp/test_smt_mm.py", line 162, in triton_mm
    mm_kernel[grid](
    ~~~~~~~~~~~~~~~^
        a,
        ^^
    ...<22 lines>...
        MICRO_K=8,
        ^^^^^^^^^^
    )
    ^
  File "/home/zhaoyy/triton-exp/.venv/lib/python3.13/site-packages/triton/runtime/jit.py", line 370, in <lambda>
    return lambda *args, **kwargs: self.run(grid=grid, warmup=False, *args, **kwargs)
                                   ~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/zhaoyy/triton-exp/.venv/lib/python3.13/site-packages/triton/runtime/jit.py", line 733, in run
    kernel = self._do_compile(key, signature, device, constexprs, options, attrs, warmup)
  File "/home/zhaoyy/triton-exp/.venv/lib/python3.13/site-packages/triton/runtime/jit.py", line 860, in _do_compile
    kernel = self.compile(src, target=target, options=options.__dict__)
  File "/home/zhaoyy/triton-exp/.venv/lib/python3.13/site-packages/triton/compiler/compiler.py", line 327, in compile
    next_module = compile_ir(module, metadata)
  File "/home/zhaoyy/triton-exp/.venv/lib/python3.13/site-packages/triton/backends/spine_triton/compiler.py", line 319, in <lambda>
    stages["llir"] = lambda src, metadata: _optimize_llir(_spine_mlir_linalgdir_to_llir(src, metadata))
                                                          ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^
  File "/home/zhaoyy/triton-exp/.venv/lib/python3.13/site-packages/triton/backends/spine_triton/compiler.py", line 93, in _spine_mlir_linalgdir_to_llir
    subprocess.check_call(
    ~~~~~~~~~~~~~~~~~~~~~^
        cmd_str,
        ^^^^^^^^
        shell=True,
        ^^^^^^^^^^^
    )
    ^
  File "/usr/lib/python3.13/subprocess.py", line 419, in check_call
    raise CalledProcessError(retcode, cmd)
subprocess.CalledProcessError: Command '/home/zhaoyy/triton-exp/.venv/lib/python3.13/site-packages/triton/backends/spine_triton/bin/spine-opt /tmp/tmpob36d_xd/linalg.mlir --spine-triton-e2e-pipeline="enable-always-tls=0 enable-fuse-group=false" -o /tmp/tmpob36d_xd/ll.mlir' returned non-zero exit status 134.

官方的 test_smt_mm.py 就会触发这个错误(把 kernel 的参数从 SPLIT_K 换成 SPLIT_M 来启用含有 parallel 的路径)。也可以用下面的简单例子复现:

import torch
import triton
import triton.language as tl
import triton.language.extra.smt as smt
from triton.backends.spine_triton.driver import CPUDriver

triton.runtime.driver.set_active(CPUDriver())

TOTAL = 128

@triton.jit
def parallel_test(x_ptr, n_elements, SUB: tl.constexpr, NSUB: tl.constexpr):
    for s in smt.parallel(0, NSUB):
        block_ptr = tl.make_block_ptr(
            base=x_ptr,
            shape=[n_elements],
            strides=[1],
            offsets=[s * SUB],
            block_shape=[SUB],
            order=[0],
        )
        v = tl.load(block_ptr, boundary_check=(0, ))
        tl.store(block_ptr, v + 1.0, boundary_check=(0, ))


def main():
    nsub = 4
    sub = TOTAL // nsub

    x = torch.zeros(TOTAL, dtype=torch.float32, device="cpu")

    parallel_test[(1, )](x, TOTAL, SUB=sub, NSUB=nsub)



if __name__ == "__main__":
    main()

请问该怎样正确使用这一接口?

环境:
python 3.13.12
triton 3.6.0+spacemit.a5