This is an automated email from the ASF dual-hosted git repository. tqchen pushed a commit to branch script/canonical-parser-df in repository https://gitbox.apache.org/repos/asf/tvm.git
commit adec73cf33545bc2129b649182cabbb16cfa9b20 Author: Tianqi Chen <[email protected]> AuthorDate: Thu Sep 24 05:16:49 2026 +0000 Match Triton calls to compiled runtime arguments Exclude constexpr values from runtime parameters and append the compiled global and profile scratch pointers in ABI order. Supply null pointers for zero-sized scratch requirements and reject unsupported scratch allocation. --- python/tvm/tirx/script/builder/triton.py | 8 +++++++- tests/python/contrib/test_tir_triton_integration.py | 14 ++++++++++---- 2 files changed, 17 insertions(+), 5 deletions(-) diff --git a/python/tvm/tirx/script/builder/triton.py b/python/tvm/tirx/script/builder/triton.py index e9038eaf33..f10c8d27dc 100644 --- a/python/tvm/tirx/script/builder/triton.py +++ b/python/tvm/tirx/script/builder/triton.py @@ -117,7 +117,6 @@ class TritonKernel(BaseKernel): if kernel_params[i].is_constexpr: constants[kernel_params[i].name] = get_const_int(arg) signature[kernel_params[i].name] = "constexpr" - kernel_args.append(arg) continue if isinstance(arg.ty, PointerType): assert isinstance(arg.ty.element_type, PrimType) @@ -133,4 +132,11 @@ class TritonKernel(BaseKernel): # TODO: Add specialization for aligned buffer pointers source = triton.compiler.ASTSource(fn=func, signature=signature, constexprs=constants) compiled = triton.compiler.compile(source, options=kwargs) + # Triton appends scratch pointers after the non-constexpr arguments. + for scratch in ("global_scratch_size", "profile_scratch_size"): + size = getattr(compiled.metadata, scratch, None) + if size is not None: + if size != 0: + raise ValueError(f"Triton kernels requiring {scratch}={size} are not supported") + kernel_args.append(tirx.reinterpret("handle", tirx.IntImm("uint64", 0))) return compiled, kernel_args diff --git a/tests/python/contrib/test_tir_triton_integration.py b/tests/python/contrib/test_tir_triton_integration.py index 4b13575408..910d730b6a 100644 --- a/tests/python/contrib/test_tir_triton_integration.py +++ b/tests/python/contrib/test_tir_triton_integration.py @@ -61,6 +61,8 @@ def test_tir_triton_integration(): output = x + y tl.store(output_ptr + offsets, output, mask=mask) + BLOCK_SIZE = 64 + @I.ir_module(s_tir=True) class Module: @T.prim_func(s_tir=True) @@ -73,7 +75,6 @@ def test_tir_triton_integration(): with T.sblock("root"): T.reads(x[0:m], y[0:m]) T.writes(output[0:m]) - BLOCK_SIZE = T.meta_var(64) T.call_kernel( add_kernel, (T.ceildiv(m, BLOCK_SIZE),), @@ -93,8 +94,13 @@ def test_tir_triton_integration(): R.output(output) return output - # Constexpr parameters (BLOCK_SIZE) stay in the kernel arguments, and the - # thread extent is 256 because the kernel is compiled with num_warps=8. + # Triton omits constexpr arguments and appends global scratch, plus profile + # scratch starting in 3.5. This kernel requires no scratch allocation. + scratch_args = [tvm.tirx.reinterpret("handle", tvm.tirx.IntImm("uint64", 0))] + if version.parse(triton.__version__) >= version.parse("3.5.0"): + scratch_args.append(tvm.tirx.reinterpret("handle", tvm.tirx.IntImm("uint64", 0))) + + # The thread extent is 256 because the kernel is compiled with num_warps=8. @I.ir_module(s_tir=True) class Parsed: @T.prim_func(s_tir=True) @@ -112,7 +118,7 @@ def test_tir_triton_integration(): y.data, output.data, m, - 64, + *scratch_args, 256, (m + T.int64(64) - T.int64(1)) // T.int64(64), )
