This is an automated email from the ASF dual-hosted git repository.
tqchen pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/main by this push:
new 4b86188893 [REFACTOR][CUDA] Store tcgen05 descriptor options in call
attributes (#20596)
4b86188893 is described below
commit 4b86188893aeb5480d5c725dd2e179fc1a487313
Author: Tianqi Chen <[email protected]>
AuthorDate: Wed Oct 7 13:34:29 2026 -0400
[REFACTOR][CUDA] Store tcgen05 descriptor options in call attributes
(#20596)
Move tcgen05 descriptor configuration into typed CUDA call attributes,
preserving the in-place destination and fixed void result. Use shared
keyword construction and reflected printing, and pass attributes to
descriptor code generation.
Remove unused scale-address operands while retaining runtime scale-ID
patching and the descriptor bit layout.
---
include/tvm/tirx/op_attr_types.h | 9 ++
python/tvm/backend/cuda/codegen/registry.py | 7 +-
python/tvm/backend/cuda/cpp/descriptors.py | 97 ++++++------------
python/tvm/backend/cuda/op.py | 82 ++++++++--------
.../cuda/tile_primitive/gemm_async/tcgen05.py | 18 ----
python/tvm/ir/op.py | 36 +++++--
python/tvm/tirx/script/ir_builder/op.py | 15 +--
src/backend/cuda/codegen/codegen_cuda.cc | 2 +-
src/backend/cuda/op/attrs.h | 108 +++++++++++++++++++++
src/backend/cuda/op/target_builtin.cc | 21 ++--
src/relax/script/printer/call.cc | 79 +++++++++++++++
src/tirx/script/printer/expr.cc | 52 ----------
.../codegen/test_tcgen05_descriptor_encoders.py | 2 -
.../cuda/gemm_async/test_gemm_async.py | 9 +-
.../python/tirx/script/test_tirx_script_printer.py | 8 +-
15 files changed, 328 insertions(+), 217 deletions(-)
diff --git a/include/tvm/tirx/op_attr_types.h b/include/tvm/tirx/op_attr_types.h
index 08f977ffdc..aab39d4ca1 100644
--- a/include/tvm/tirx/op_attr_types.h
+++ b/include/tvm/tirx/op_attr_types.h
@@ -75,6 +75,15 @@ using FLegalize = ffi::TypedFunction<PrimExpr(PrimExpr)>;
*/
using TScriptPrinterName = ffi::String;
+/*!
+ * \brief The published script callable accepts the shared Op construction
contract.
+ *
+ * Published by Python construction namespaces, not by Op definitions: a
printer
+ * name alone does not establish a legacy wrapper's positional/keyword
contract.
+ * Semantic syntax hooks retain precedence over ordinary named construction.
+ */
+using TScriptStandardCall = bool;
+
/*!
* \brief Specifies that TVMScript printer prints the dtype as the first/last
argument.
If not specified, dtype will not be printed.
diff --git a/python/tvm/backend/cuda/codegen/registry.py
b/python/tvm/backend/cuda/codegen/registry.py
index 7ae1235185..94acd6ff18 100644
--- a/python/tvm/backend/cuda/codegen/registry.py
+++ b/python/tvm/backend/cuda/codegen/registry.py
@@ -66,8 +66,11 @@ def register_codegen(op, backend="cuda"):
op_names = {full_op_name, canonical_op_name}
@functools.wraps(func)
- def wrapper(arg_list):
- res = func(*arg_list) # pylint: disable=not-callable
+ def wrapper(arg_list, attrs=None):
+ # Keep operand-only helpers callable directly while forwarding the
+ # typed static options of an attributed Call to its lowering.
+ kwargs = {} if attrs is None else {"attrs": attrs}
+ res = func(*arg_list, **kwargs) # pylint: disable=not-callable
if isinstance(res, tuple):
return res[0], res[1]
return res, list()
diff --git a/python/tvm/backend/cuda/cpp/descriptors.py
b/python/tvm/backend/cuda/cpp/descriptors.py
index bf30b54e0b..3e246d1e07 100644
--- a/python/tvm/backend/cuda/cpp/descriptors.py
+++ b/python/tvm/backend/cuda/cpp/descriptors.py
@@ -30,7 +30,7 @@ same constraints into its dispatch.
from ..codegen.registry import CODEGEN_REGISTRY, register_codegen
from ..codegen.schema import device_intrinsic
from ..codegen.types import PTXDataType
-from ..codegen.utils import parse_str, validate_cta_group
+from ..codegen.utils import validate_cta_group
# =============================================================================
# wgmma_encode_matrix_descriptor — pure-C bitfield struct fill (no asm).
@@ -321,37 +321,21 @@ device_intrinsic(
@register_codegen("cuda_tcgen05_encode_instr_descriptor")
-def codegen_cuda_tcgen05_encode_instr_descriptor(
- desc,
- d_dtype,
- a_dtype,
- b_dtype,
- M,
- N,
- K,
- trans_a,
- trans_b,
- n_cta_group,
- neg_a,
- neg_b,
- sat_d,
- is_sparse,
-):
- """Validate dtype combinations and shape, translate dtypes to PTX format
- integers, then forward to the schema-driven impl."""
- a_dtype = parse_str(a_dtype)
- b_dtype = parse_str(b_dtype)
- d_dtype = parse_str(d_dtype)
- M = int(M)
- N = int(N)
- K = int(K)
- n_cta_group = validate_cta_group(n_cta_group)
- trans_a = bool(trans_a)
- trans_b = bool(trans_b)
- neg_a = bool(neg_a)
- neg_b = bool(neg_b)
- sat_d = bool(sat_d)
- is_sparse = bool(is_sparse)
+def codegen_cuda_tcgen05_encode_instr_descriptor(desc, *, attrs):
+ """Validate static Call.attrs and lower the in-place descriptor fill."""
+ d_dtype = attrs.d_dtype
+ a_dtype = attrs.a_dtype
+ b_dtype = attrs.b_dtype
+ M = attrs.M
+ N = attrs.N
+ K = attrs.K
+ trans_a = attrs.trans_a
+ trans_b = attrs.trans_b
+ n_cta_group = validate_cta_group(attrs.n_cta_groups)
+ neg_a = attrs.neg_a
+ neg_b = attrs.neg_b
+ sat_d = attrs.sat_d
+ is_sparse = attrs.is_sparse
kind = _get_tcgen05_mma_kind(d_dtype, a_dtype, b_dtype)
if kind not in ["f16", "tf32", "f8f6f4", "i8"]:
@@ -436,39 +420,22 @@ device_intrinsic(
@register_codegen("cuda_tcgen05_encode_instr_descriptor_block_scaled")
-def codegen_cuda_tcgen05_encode_instr_descriptor_block_scaled(
- desc,
- d_dtype,
- a_dtype,
- b_dtype,
- sfa_dtype,
- sfb_dtype,
- sfa_tmem_addr,
- sfb_tmem_addr,
- M,
- N,
- K,
- trans_a,
- trans_b,
- n_cta_group,
- neg_a,
- neg_b,
- is_sparse,
-):
- a_dtype = parse_str(a_dtype)
- b_dtype = parse_str(b_dtype)
- d_dtype = parse_str(d_dtype)
- sfa_dtype = parse_str(sfa_dtype)
- sfb_dtype = parse_str(sfb_dtype)
- M = int(M)
- N = int(N)
- K = int(K)
- n_cta_group = validate_cta_group(n_cta_group)
- trans_a = bool(trans_a)
- trans_b = bool(trans_b)
- neg_a = bool(neg_a)
- neg_b = bool(neg_b)
- is_sparse = bool(is_sparse)
+def codegen_cuda_tcgen05_encode_instr_descriptor_block_scaled(desc, *, attrs):
+ """Validate static Call.attrs and lower the in-place descriptor fill."""
+ d_dtype = attrs.d_dtype
+ a_dtype = attrs.a_dtype
+ b_dtype = attrs.b_dtype
+ sfa_dtype = attrs.sfa_dtype
+ sfb_dtype = attrs.sfb_dtype
+ M = attrs.M
+ N = attrs.N
+ K = attrs.K
+ trans_a = attrs.trans_a
+ trans_b = attrs.trans_b
+ n_cta_group = validate_cta_group(attrs.n_cta_groups)
+ neg_a = attrs.neg_a
+ neg_b = attrs.neg_b
+ is_sparse = attrs.is_sparse
kind = _get_tcgen05_mma_kind(d_dtype, a_dtype, b_dtype, sfa_dtype,
sfb_dtype)
valid_kinds = {"mxf8f6f4", "mxf4", "mxf4nvf4"}
diff --git a/python/tvm/backend/cuda/op.py b/python/tvm/backend/cuda/op.py
index 63793ebf50..4218095fc8 100644
--- a/python/tvm/backend/cuda/op.py
+++ b/python/tvm/backend/cuda/op.py
@@ -21,7 +21,7 @@ from __future__ import annotations
from tvm import tirx
from tvm.ir import Call, Op, StringImm
-from tvm.ir.op import _init_op_api
+from tvm.ir.op import _init_op_api, _make_op_api
from tvm.ir.type import PointerType, PrimType
from tvm.runtime import const
from tvm.tirx.op import bitwise_and, call_intrin, tvm_access_ptr
@@ -603,6 +603,14 @@ def ptx_legacy_ldmatrix(*all_args):
)
+_encode_instr_descriptor = _make_op_api(
+ Op.get("tirx.cuda.tcgen05_encode_instr_descriptor"), __name__
+)
+_encode_instr_descriptor_block_scaled = _make_op_api(
+ Op.get("tirx.cuda.tcgen05_encode_instr_descriptor_block_scaled"), __name__
+)
+
+
def cuda_tcgen05_encode_instr_descriptor(
desc,
*,
@@ -669,23 +677,21 @@ def cuda_tcgen05_encode_instr_descriptor(
Whether the MMA operation is sparse.
"""
_choice("n_cta_groups", n_cta_groups, _TCGEN05_CTA_GROUP)
- return call_intrin(
- "",
- "tirx.cuda.tcgen05_encode_instr_descriptor",
+ return _encode_instr_descriptor(
desc,
- d_dtype,
- a_dtype,
- b_dtype,
- M,
- N,
- K,
- trans_a,
- trans_b,
- n_cta_groups,
- neg_a,
- neg_b,
- sat_d,
- is_sparse,
+ d_dtype=d_dtype,
+ a_dtype=a_dtype,
+ b_dtype=b_dtype,
+ M=M,
+ N=N,
+ K=K,
+ trans_a=trans_a,
+ trans_b=trans_b,
+ n_cta_groups=n_cta_groups,
+ neg_a=neg_a,
+ neg_b=neg_b,
+ sat_d=sat_d,
+ is_sparse=is_sparse,
)
@@ -697,8 +703,6 @@ def cuda_tcgen05_encode_instr_descriptor_block_scaled(
b_dtype,
sfa_dtype,
sfb_dtype,
- sfa_tmem_addr,
- sfb_tmem_addr,
M,
N,
K,
@@ -731,12 +735,6 @@ def cuda_tcgen05_encode_instr_descriptor_block_scaled(
sfb_dtype : str
The datatype of scale factor matrix B.
- sfa_tmem_addr : Expr
- The address of the scale factor matrix A in tensor memory, should be
uint32_t.
-
- sfb_tmem_addr : Expr
- The address of the scale factor matrix B in tensor memory, should be
uint32_t.
-
M : int
The size of non-reduction dimension of Matrix A.
@@ -767,26 +765,22 @@ def cuda_tcgen05_encode_instr_descriptor_block_scaled(
Whether the MMA operation is sparse.
"""
_choice("n_cta_groups", n_cta_groups, _TCGEN05_CTA_GROUP)
- return call_intrin(
- "",
- "tirx.cuda.tcgen05_encode_instr_descriptor_block_scaled",
+ return _encode_instr_descriptor_block_scaled(
desc,
- d_dtype,
- a_dtype,
- b_dtype,
- sfa_dtype,
- sfb_dtype,
- sfa_tmem_addr,
- sfb_tmem_addr,
- M,
- N,
- K,
- trans_a,
- trans_b,
- n_cta_groups,
- neg_a,
- neg_b,
- is_sparse,
+ d_dtype=d_dtype,
+ a_dtype=a_dtype,
+ b_dtype=b_dtype,
+ sfa_dtype=sfa_dtype,
+ sfb_dtype=sfb_dtype,
+ M=M,
+ N=N,
+ K=K,
+ trans_a=trans_a,
+ trans_b=trans_b,
+ n_cta_groups=n_cta_groups,
+ neg_a=neg_a,
+ neg_b=neg_b,
+ is_sparse=is_sparse,
)
diff --git a/python/tvm/backend/cuda/tile_primitive/gemm_async/tcgen05.py
b/python/tvm/backend/cuda/tile_primitive/gemm_async/tcgen05.py
index acb3b7c90c..f3e43707bc 100644
--- a/python/tvm/backend/cuda/tile_primitive/gemm_async/tcgen05.py
+++ b/python/tvm/backend/cuda/tile_primitive/gemm_async/tcgen05.py
@@ -1279,23 +1279,6 @@ def gemm_async_tcgen05_impl(op_call: TilePrimitiveCall,
sctx: DispatchContext) -
sfa_base = SFA_buffer.allocated_addr[0]
sfb_base = SFB_buffer.allocated_addr[0]
- # Compute initial SFA/SFB addresses (for ki=0). Both physical axes are
- # part of a TMEM address: TLane occupies the high half-word and TCol
- # the low half-word. Dropping TLane silently redirects an explicitly
- # banded scale view to row zero.
- sfa_coord_0 = SFA_slice_layout.apply(0)
- sfb_coord_0 = SFB_slice_layout.apply(0)
- sfa_tlane_0 = sfa_coord_0.get("TLane", 0)
- sfb_tlane_0 = sfb_coord_0.get("TLane", 0)
- sfa_tcol_0 = sfa_coord_0.get("TCol", 0)
- sfb_tcol_0 = sfb_coord_0.get("TCol", 0)
- SFA_init_addr = _get_tmem_addr_fast(
- sfa_base, sfa_tlane_0, tvm.tirx.floordiv(sfa_tcol_0,
SFA_elem_per_col)
- )
- SFB_init_addr = _get_tmem_addr_fast(
- sfb_base, sfb_tlane_0, tvm.tirx.floordiv(sfb_tcol_0,
SFB_elem_per_col)
- )
-
# Rotate sf_id per ki when multiple ki share one SF column.
needs_sf_id = sfa_sf_mma_k < SFA_elem_per_col and sfa_elems_per_ki > 0
@@ -1503,7 +1486,6 @@ def gemm_async_tcgen05_impl(op_call: TilePrimitiveCall,
sctx: DispatchContext) -
def impl():
descI_local: T.uint32
T.cuda.tcgen05.encode_instr_descriptor_block_scaled(T.address_of(descI_local),
d_dtype=C_type, a_dtype=A_type, b_dtype=B_type, sfa_dtype=SFA_type,
sfb_dtype=SFB_type, # noqa: E501, F821
-
sfa_tmem_addr=SFA_init_addr, sfb_tmem_addr=SFB_init_addr, # noqa: E501
M=M_mma *
cta_group, N=N_mma, K=MMA_K, trans_a=a_mn_major, trans_b=b_mn_major,
n_cta_groups=cta_group) # noqa: E501
call_main(descI_local) # noqa: F821
else:
diff --git a/python/tvm/ir/op.py b/python/tvm/ir/op.py
index 927b8edac7..0b406b556f 100644
--- a/python/tvm/ir/op.py
+++ b/python/tvm/ir/op.py
@@ -27,12 +27,23 @@ import tvm_ffi
from . import _ffi_api
from .expr import Expr
+_RESULT_TYPE_UNSET = object()
+
def _make_op_api(op, module_name):
"""Build a callable whose operands and result are governed by its Op."""
+ from .attrs import make_node # pylint: disable=import-outside-toplevel
from .expr import Call, reinfer_type # pylint:
disable=import-outside-toplevel
- def call(*args, attrs=None, ty_args=None, span=None, ret_ty=None,
**kwargs):
+ def call(
+ *args,
+ attrs=None,
+ ty_args=None,
+ span=None,
+ ty=_RESULT_TYPE_UNSET,
+ ret_ty=_RESULT_TYPE_UNSET,
+ **kwargs,
+ ):
# Bind named operands using the registered signature, without inventing
# defaults or interpreting arbitrary Python wrapper signatures.
operands = list(args)
@@ -45,18 +56,29 @@ def _make_op_api(op, module_name):
else:
raise TypeError(f"{op.name}: missing operand {info.name!r}")
if kwargs:
- raise TypeError(f"{op.name}: unexpected keyword operands
{tuple(kwargs)}")
- if ret_ty is None and (
+ if not op.attrs_type_key:
+ raise TypeError(f"{op.name}: unexpected keyword operands
{tuple(kwargs)}")
+ if attrs is not None:
+ raise TypeError(f"{op.name}: cannot mix attrs with attribute
keywords")
+ attrs = make_node(op.attrs_type_key, **kwargs)
+ if ty is not _RESULT_TYPE_UNSET:
+ if ret_ty is not _RESULT_TYPE_UNSET:
+ raise TypeError(f"{op.name}: ty and ret_ty cannot both be
specified")
+ ret_ty = ty
+ elif (ret_ty is _RESULT_TYPE_UNSET or ret_ty is None) and (
op.get_attr("TFixedReturnType") is not None or
op.get_attr("FInferType") is not None
):
provisional = Call.unchecked(op, operands, attrs=attrs,
ty_args=ty_args, span=span)
ret_ty = reinfer_type(provisional)
+ if ret_ty is _RESULT_TYPE_UNSET:
+ ret_ty = None
return Call(op, operands, attrs=attrs, ty_args=ty_args, span=span,
ty=ret_ty)
call.__name__ = op.name.rsplit(".", 1)[-1]
call.__module__ = module_name
call.__doc__ = op.doc or f"Construct a call to {op.name}."
call.__tvm_op__ = op
+ call.__tvm_standard_call__ = True
return call
@@ -69,9 +91,11 @@ def _init_op_api(namespace, target_module_name=None):
module or SimpleNamespace containers; underscores are ordinary name parts.
Generated functions accept registered positional/named operands plus
- ``attrs``, ``ty_args``, ``span`` and ``ret_ty``. Omitting the result
invokes
- an available Op inference hook; without one, Call retains a missing type.
- Explicit results and inference/validation errors are preserved.
+ ``attrs``, ``ty_args``, ``span`` and ``ty`` (also spelled ``ret_ty``).
+ Omitting the result invokes an available Op inference hook; without one,
+ Call retains a missing type.
+ Explicit results and inference/validation errors are preserved. Attribute
+ keywords construct the registered attrs schema, when one is declared.
Existing generated functions are reused. A deliberate wrapper may declare
``__tvm_op__ = Op.get(name)`` to retain ownership of that name. This
declares
diff --git a/python/tvm/tirx/script/ir_builder/op.py
b/python/tvm/tirx/script/ir_builder/op.py
index 8bd107da37..3d189e7281 100644
--- a/python/tvm/tirx/script/ir_builder/op.py
+++ b/python/tvm/tirx/script/ir_builder/op.py
@@ -335,13 +335,14 @@ def register_script_namespace(
current = getattr(current, part, None)
op = _ir.Op.get(op_name)
identity = getattr(current, "__tvm_op__", None)
- if (
- callable(current)
- and isinstance(identity, _ir.Op)
- and identity.same_as(op)
- and op.get_attr("TScriptPrinterName") is None
- ):
- op.set_attr("TScriptPrinterName", op_name)
+ if callable(current) and isinstance(identity, _ir.Op) and
identity.same_as(op):
+ op.set_attr(
+ "TScriptStandardCall",
+ bool(getattr(current, "__tvm_standard_call__", False)),
+ override=True,
+ )
+ if op.get_attr("TScriptPrinterName") is None:
+ op.set_attr("TScriptPrinterName", op_name)
else:
_register_script_namespace_printer_names(namespace, f"tirx.{name}",
override)
return namespace
diff --git a/src/backend/cuda/codegen/codegen_cuda.cc
b/src/backend/cuda/codegen/codegen_cuda.cc
index 66598b90ab..fbfe880efd 100644
--- a/src/backend/cuda/codegen/codegen_cuda.cc
+++ b/src/backend/cuda/codegen/codegen_cuda.cc
@@ -1065,7 +1065,7 @@ void CodeGenCUDA::Dispatch_(const CallNode* op,
std::ostream& os) {
auto codegen =
codegen_getter.value()(call_op->name).cast<ffi::Optional<tvm::ffi::Function>>();
if (codegen.has_value()) {
// codegen is registered, it should return a Call to cuda_func_call
- auto func_call = codegen.value()(op->args);
+ auto func_call = codegen.value()(op->args, op->attrs);
auto res = func_call.cast<ffi::Tuple<Call, ffi::Array<ffi::String>>>();
print_cuda_func_call(res.get<0>().get(), os);
for (const auto& tag : res.get<1>()) {
diff --git a/src/backend/cuda/op/attrs.h b/src/backend/cuda/op/attrs.h
new file mode 100644
index 0000000000..93fa8d3fa3
--- /dev/null
+++ b/src/backend/cuda/op/attrs.h
@@ -0,0 +1,108 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+#ifndef TVM_BACKEND_CUDA_OP_ATTRS_H_
+#define TVM_BACKEND_CUDA_OP_ATTRS_H_
+
+#include <tvm/ir/attrs.h>
+
+namespace tvm {
+namespace tirx {
+
+/*! \brief Static options for the dense tcgen05 instruction descriptor. */
+struct TCGen05InstrDescriptorAttrs : public AttrsNode {
+ ffi::String d_dtype;
+ ffi::String a_dtype;
+ ffi::String b_dtype;
+ int64_t M;
+ int64_t N;
+ int64_t K;
+ bool trans_a;
+ bool trans_b;
+ int64_t n_cta_groups = 1;
+ bool neg_a = false;
+ bool neg_b = false;
+ bool sat_d = false;
+ bool is_sparse = false;
+
+ static void RegisterReflection() {
+ namespace refl = ffi::reflection;
+ refl::ObjectDef<TCGen05InstrDescriptorAttrs>()
+ .def_ro("d_dtype", &TCGen05InstrDescriptorAttrs::d_dtype)
+ .def_ro("a_dtype", &TCGen05InstrDescriptorAttrs::a_dtype)
+ .def_ro("b_dtype", &TCGen05InstrDescriptorAttrs::b_dtype)
+ .def_ro("M", &TCGen05InstrDescriptorAttrs::M)
+ .def_ro("N", &TCGen05InstrDescriptorAttrs::N)
+ .def_ro("K", &TCGen05InstrDescriptorAttrs::K)
+ .def_ro("trans_a", &TCGen05InstrDescriptorAttrs::trans_a)
+ .def_ro("trans_b", &TCGen05InstrDescriptorAttrs::trans_b)
+ .def_ro("n_cta_groups", &TCGen05InstrDescriptorAttrs::n_cta_groups,
refl::DefaultValue(1))
+ .def_ro("neg_a", &TCGen05InstrDescriptorAttrs::neg_a,
refl::DefaultValue(false))
+ .def_ro("neg_b", &TCGen05InstrDescriptorAttrs::neg_b,
refl::DefaultValue(false))
+ .def_ro("sat_d", &TCGen05InstrDescriptorAttrs::sat_d,
refl::DefaultValue(false))
+ .def_ro("is_sparse", &TCGen05InstrDescriptorAttrs::is_sparse,
refl::DefaultValue(false));
+ }
+ TVM_FFI_DECLARE_OBJECT_INFO_FINAL("tirx.cuda.TCGen05InstrDescriptorAttrs",
+ TCGen05InstrDescriptorAttrs, AttrsNode);
+};
+
+/*! \brief Static options for the block-scaled tcgen05 instruction descriptor.
*/
+struct TCGen05InstrDescriptorBlockScaledAttrs : public AttrsNode {
+ ffi::String d_dtype;
+ ffi::String a_dtype;
+ ffi::String b_dtype;
+ ffi::String sfa_dtype;
+ ffi::String sfb_dtype;
+ int64_t M;
+ int64_t N;
+ int64_t K;
+ bool trans_a;
+ bool trans_b;
+ int64_t n_cta_groups = 1;
+ bool neg_a = false;
+ bool neg_b = false;
+ bool is_sparse = false;
+
+ static void RegisterReflection() {
+ namespace refl = ffi::reflection;
+ refl::ObjectDef<TCGen05InstrDescriptorBlockScaledAttrs>()
+ .def_ro("d_dtype", &TCGen05InstrDescriptorBlockScaledAttrs::d_dtype)
+ .def_ro("a_dtype", &TCGen05InstrDescriptorBlockScaledAttrs::a_dtype)
+ .def_ro("b_dtype", &TCGen05InstrDescriptorBlockScaledAttrs::b_dtype)
+ .def_ro("sfa_dtype",
&TCGen05InstrDescriptorBlockScaledAttrs::sfa_dtype)
+ .def_ro("sfb_dtype",
&TCGen05InstrDescriptorBlockScaledAttrs::sfb_dtype)
+ .def_ro("M", &TCGen05InstrDescriptorBlockScaledAttrs::M)
+ .def_ro("N", &TCGen05InstrDescriptorBlockScaledAttrs::N)
+ .def_ro("K", &TCGen05InstrDescriptorBlockScaledAttrs::K)
+ .def_ro("trans_a", &TCGen05InstrDescriptorBlockScaledAttrs::trans_a)
+ .def_ro("trans_b", &TCGen05InstrDescriptorBlockScaledAttrs::trans_b)
+ .def_ro("n_cta_groups",
&TCGen05InstrDescriptorBlockScaledAttrs::n_cta_groups,
+ refl::DefaultValue(1))
+ .def_ro("neg_a", &TCGen05InstrDescriptorBlockScaledAttrs::neg_a,
refl::DefaultValue(false))
+ .def_ro("neg_b", &TCGen05InstrDescriptorBlockScaledAttrs::neg_b,
refl::DefaultValue(false))
+ .def_ro("is_sparse",
&TCGen05InstrDescriptorBlockScaledAttrs::is_sparse,
+ refl::DefaultValue(false));
+ }
+
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("tirx.cuda.TCGen05InstrDescriptorBlockScaledAttrs",
+ TCGen05InstrDescriptorBlockScaledAttrs,
AttrsNode);
+};
+
+} // namespace tirx
+} // namespace tvm
+
+#endif // TVM_BACKEND_CUDA_OP_ATTRS_H_
diff --git a/src/backend/cuda/op/target_builtin.cc
b/src/backend/cuda/op/target_builtin.cc
index 0306dfb38f..2ef90b9f13 100644
--- a/src/backend/cuda/op/target_builtin.cc
+++ b/src/backend/cuda/op/target_builtin.cc
@@ -31,6 +31,8 @@
#include <string>
+#include "attrs.h"
+
namespace tvm {
namespace tirx {
namespace builtin {
@@ -58,6 +60,9 @@ void RegisterCudaTargetBuiltins() {
if (registered) return;
registered = true;
+ TCGen05InstrDescriptorAttrs::RegisterReflection();
+ TCGen05InstrDescriptorBlockScaledAttrs::RegisterReflection();
+
OpDef("tirx.tvm_load_matrix_sync")
.signature(
sig::arg("fragment", "The matrix fragment."), sig::arg<IntExpr>("m",
"The M dimension."),
@@ -407,20 +412,12 @@ void RegisterDeviceIntrinsics() {
sig::arg("cond"))
.set_attr<TFixedReturnType>("TFixedReturnType", PrimType::Int(64));
RegisterDeviceIntrinsic(OpDef("tirx.cuda.tcgen05_encode_instr_descriptor"),
"cuda",
- CallEffectKind::kOpaque, sig::arg("desc"),
sig::arg("d_dtype"),
- sig::arg("a_dtype"), sig::arg("b_dtype"),
sig::arg<IntExpr>("M"),
- sig::arg<IntExpr>("N"), sig::arg<IntExpr>("K"),
sig::arg("trans_a"),
- sig::arg("trans_b"),
sig::arg<IntExpr>("n_cta_groups"), sig::arg("neg_a"),
- sig::arg("neg_b"), sig::arg("sat_d"),
sig::arg("is_sparse"))
+ CallEffectKind::kOpaque, sig::arg("desc"),
+ sig::call_attrs<TCGen05InstrDescriptorAttrs>())
.set_attr<TFixedReturnType>("TFixedReturnType", PrimType::Void());
RegisterDeviceIntrinsic(OpDef("tirx.cuda.tcgen05_encode_instr_descriptor_block_scaled"),
"cuda",
- CallEffectKind::kOpaque, sig::arg("desc"),
sig::arg("d_dtype"),
- sig::arg("a_dtype"), sig::arg("b_dtype"),
sig::arg("sfa_dtype"),
- sig::arg("sfb_dtype"), sig::arg("sfa_tmem_addr"),
- sig::arg("sfb_tmem_addr"), sig::arg<IntExpr>("M"),
sig::arg<IntExpr>("N"),
- sig::arg<IntExpr>("K"), sig::arg("trans_a"),
sig::arg("trans_b"),
- sig::arg<IntExpr>("n_cta_groups"),
sig::arg("neg_a"), sig::arg("neg_b"),
- sig::arg("is_sparse"))
+ CallEffectKind::kOpaque, sig::arg("desc"),
+
sig::call_attrs<TCGen05InstrDescriptorBlockScaledAttrs>())
.set_attr<TFixedReturnType>("TFixedReturnType", PrimType::Void());
RegisterDeviceIntrinsic(OpDef("tirx.cuda.tcgen05_encode_matrix_descriptor"),
"cuda",
CallEffectKind::kOpaque, sig::arg("desc"),
sig::arg("addr"),
diff --git a/src/relax/script/printer/call.cc b/src/relax/script/printer/call.cc
index a9f1f6a812..103757c025 100644
--- a/src/relax/script/printer/call.cc
+++ b/src/relax/script/printer/call.cc
@@ -242,10 +242,89 @@ TVM_FFI_STATIC_INIT_BLOCK() {
}
}
+// The name is published only for a callable implementing the shared Op
contract.
+// Legacy/semantic wrappers retain their existing translators below.
+ffi::Optional<ExprDoc> StandardCallDocTranslate(DocTranslatorObj* d, const
CallNode* call) {
+ auto maybe_op = call->op.as<Op>();
+ if (!maybe_op || !Op::HasAttrMap("TScriptStandardCall")) return std::nullopt;
+ const Op& op = maybe_op.value();
+ if (op->attrs_type_key.empty()) return std::nullopt;
+ static const auto& standard =
Op::GetAttrMap<tirx::TScriptStandardCall>("TScriptStandardCall");
+ static const auto& names =
Op::GetAttrMap<tirx::TScriptPrinterName>("TScriptPrinterName");
+ if (!standard.get(op, false) || !names.count(op) || names[op].empty())
return std::nullopt;
+ if (call->attrs.defined() ? op->attrs_type_key != call->attrs->GetTypeKey()
+ : !op->attrs_type_key.empty()) {
+ return RawCall(d, call);
+ }
+ try {
+ op.Validate(call);
+ } catch (const ffi::Error&) {
+ return RawCall(d, call);
+ }
+ ffi::Array<ExprDoc> args;
+ for (const Expr& arg : call->args) {
+ args.push_back(MaterializeCallArgument(d, arg, d->Translate(arg).value()));
+ }
+ ffi::Array<ffi::String> keys;
+ ffi::Array<ExprDoc> values;
+ if (!call->ty_args.empty()) {
+ ffi::Array<ExprDoc> types;
+ for (const Type& type : call->ty_args) types.push_back(TypeValue(d, type,
false));
+ keys.push_back("ty_args");
+ values.push_back(ListDoc(types));
+ }
+ if (call->attrs.defined()) {
+ std::vector<std::pair<ffi::String, ffi::Any>> fields;
+ ffi::reflection::ForEachFieldInfo(
+ TVMFFIGetTypeInfo(call->attrs->type_index()), [&](const
TVMFFIFieldInfo* field) {
+ ffi::Any value = ffi::reflection::FieldGetter(field)(call->attrs);
+ // Factory defaults can have effects or change between constructions.
+ // Only a reflected literal default is safe to omit here.
+ if ((field->flags & kTVMFFIFieldFlagBitMaskHasDefault) &&
+ !(field->flags & kTVMFFIFieldFlagBitMaskDefaultFromFactory) &&
+ ffi::StructuralEqual()(
+ value,
ffi::AnyView::CopyFromTVMFFIAny(field->default_value_or_factory))) {
+ return;
+ }
+ fields.emplace_back(ffi::String(field->name), std::move(value));
+ });
+ // Without attribute keywords the named builder does not construct attrs.
+ // Preserve even an empty or entirely default-valued schema explicitly.
+ if (fields.empty()) return RawCall(d, call);
+ std::sort(fields.begin(), fields.end(),
+ [](const auto& a, const auto& b) { return a.first < b.first; });
+ for (const auto& [key, value] : fields) {
+ if (key == "ty" || key == "ret_ty" || key == "attrs" || key == "ty_args"
|| key == "span" ||
+ key == "type_key" ||
+ std::any_of(op->args_info.begin(), op->args_info.end(),
+ [&](const ArgumentInfo& info) { return info->name ==
key; })) {
+ return RawCall(d, call);
+ }
+ keys.push_back(key);
+ values.push_back(AnyValue(d, value));
+ }
+ }
+ bool omit_result = false;
+ try {
+ // Match construction without ty, rather than let inference read the
stored result.
+ Call provisional =
+ Call::Unchecked(Type::Missing(), call->op, call->args, call->attrs,
call->ty_args);
+ omit_result = ffi::StructuralEqual()(Call::ReinferType(provisional.get()),
call->ty);
+ } catch (const ffi::Error&) {
+ // Missing information, unavailable inference, or errors require the exact
stored type.
+ }
+ if (!omit_result) {
+ keys.push_back("ty");
+ values.push_back(TypeValue(d, call->ty));
+ }
+ return NamedCallCallee(names[op])->Call(args, keys, values);
+}
+
ffi::Optional<ExprDoc> CallDefaultDocTranslate(DocTranslatorObj* d,
ffi::AnyView input,
const ffi::Object* destination)
{
const auto* call =
ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const
CallNode>(input);
+ if (auto doc = StandardCallDocTranslate(d, call)) return doc;
if (auto op = call->op.as<Op>(); op && op.value()->name.find("relax.") == 0)
{
if (HasRelaxCallResult(call, destination)) {
ffi::Array<ExprDoc> args;
diff --git a/src/tirx/script/printer/expr.cc b/src/tirx/script/printer/expr.cc
index 15a16e02c2..5403f771b7 100644
--- a/src/tirx/script/printer/expr.cc
+++ b/src/tirx/script/printer/expr.cc
@@ -212,58 +212,6 @@ TVM_FFI_STATIC_INIT_BLOCK() {
FDocTranslate::FromNative<&CUDAFuncCallDocTranslate>());
}
-template <bool block_scaled>
-ffi::Optional<ExprDoc> CUDAInstructionDescriptorDocTranslate(DocTranslatorObj*
d,
- ffi::AnyView
input,
- const
ffi::Object*) {
- const auto* call =
- ffi::details::AnyUnsafe::RawObjectPtrFromAnyViewAfterCheck<const
CallNode>(input);
- if (!CanTranslateExplicitResultCall(call) || call->args.size() !=
(block_scaled ? 17 : 14) ||
- !ffi::StructuralEqual()(call->ty, PrimType::Void())) {
- return RawCall(d, call);
- }
- Op op = call->op.as_or_throw<Op>();
- if (op->args_info.size() != call->args.size()) return RawCall(d, call);
- constexpr size_t optional_begin = block_scaled ? 13 : 9;
- ffi::Array<ffi::String> keys;
- ffi::Array<ExprDoc> values;
- ExprDoc descriptor = d->Translate(call->args[0]).value();
- for (size_t i = 1; i < call->args.size(); ++i) {
- const Expr& arg = call->args[i];
- if (i >= optional_begin) {
- const auto* value = arg.as<IntImmNode>();
- PrimType default_type = i == optional_begin ? PrimType::Int(32) :
PrimType::Bool();
- int64_t default_value = i == optional_begin ? 1 : 0;
- if (value && value->value == default_value &&
- ffi::StructuralEqual()(value->ty, default_type)) {
- continue;
- }
- }
- const auto* string = arg.as<StringImmNode>();
- ExprDoc value = string ? LiteralDoc::Str(string->value, std::nullopt)
- : MaterializeCallArgument(d, arg,
d->Translate(arg).value());
- d->RecordOrigin(value, arg);
- keys.push_back(op->args_info[i]->name);
- values.push_back(value);
- }
- return NamespaceDoc("tirx")
- ->Attr("cuda")
- ->Attr("tcgen05")
- ->Attr(block_scaled ? "encode_instr_descriptor_block_scaled" :
"encode_instr_descriptor")
- ->Call({descriptor}, keys, values);
-}
-
-TVM_FFI_STATIC_INIT_BLOCK() {
- OpDef("tirx.cuda.tcgen05_encode_instr_descriptor")
- .set_attr<FDocTranslate>(
- kOpCallDocTranslate,
-
FDocTranslate::FromNative<&CUDAInstructionDescriptorDocTranslate<false>>());
- OpDef("tirx.cuda.tcgen05_encode_instr_descriptor_block_scaled")
- .set_attr<FDocTranslate>(
- kOpCallDocTranslate,
-
FDocTranslate::FromNative<&CUDAInstructionDescriptorDocTranslate<true>>());
-}
-
ffi::Optional<ExprDoc> LLVMIntrinsicDocTranslate(DocTranslatorObj* d,
ffi::AnyView input,
const ffi::Object*) {
const auto* call =
diff --git a/tests/python/tirx/codegen/test_tcgen05_descriptor_encoders.py
b/tests/python/tirx/codegen/test_tcgen05_descriptor_encoders.py
index cd80b898f8..1162b143c8 100644
--- a/tests/python/tirx/codegen/test_tcgen05_descriptor_encoders.py
+++ b/tests/python/tirx/codegen/test_tcgen05_descriptor_encoders.py
@@ -125,8 +125,6 @@ def
test_instr_descriptor_block_scaled_matches_runtime_encoder(
b_dtype=b_dtype,
sfa_dtype="float8_e8m0fnu",
sfb_dtype="float8_e8m0fnu",
- sfa_tmem_addr=0,
- sfb_tmem_addr=0,
M=m,
N=n,
K=k,
diff --git
a/tests/python/tirx/operator/tile_primitive/cuda/gemm_async/test_gemm_async.py
b/tests/python/tirx/operator/tile_primitive/cuda/gemm_async/test_gemm_async.py
index ac427a0dad..0215926190 100644
---
a/tests/python/tirx/operator/tile_primitive/cuda/gemm_async/test_gemm_async.py
+++
b/tests/python/tirx/operator/tile_primitive/cuda/gemm_async/test_gemm_async.py
@@ -3834,8 +3834,13 @@ def
test_gemm_tcgen05_preserves_block_scale_tmem_lane_bases():
sctx = DispatchContext(target, ExecScope("thread"), {}, {},
scope_kind="thread")
script = gemm_async_tcgen05_impl(call, sctx).script()
- assert "T.cuda.get_tmem_addr(T.uint32(256), 1, 0)" in script
- assert "T.cuda.get_tmem_addr(T.uint32(320), 2, 0)" in script
+ # Check the addresses consumed by MMA, rather than the former unused
+ # initial-address operands of the instruction descriptor encoder.
+ assert "sfa_linear % 128 // 4 + 1" in script
+ assert "sfb_linear % 128 // 4 + 2" in script
+ assert "T.cuda.get_tmem_addr(T.uint32(256), sfa_tlane, sfa_tcol // 4)" in
script
+ assert "T.cuda.get_tmem_addr(T.uint32(320), sfb_tlane, sfb_tcol // 4)" in
script
+ assert "descI_local, sfa_addr, sfb_addr, should_accum" in script
@pytest.mark.parametrize(
diff --git a/tests/python/tirx/script/test_tirx_script_printer.py
b/tests/python/tirx/script/test_tirx_script_printer.py
index 3504ed03c0..3e541832f3 100644
--- a/tests/python/tirx/script/test_tirx_script_printer.py
+++ b/tests/python/tirx/script/test_tirx_script_printer.py
@@ -785,7 +785,7 @@ def test_printer_ptx_more():
sat_d=False,
is_sparse=False,
),
- 'd: T.handle = T.handle()\nT.cuda.tcgen05.encode_instr_descriptor(d,
d_dtype="f16", a_dtype="f16", b_dtype="f16", M=16, N=16, K=16,
trans_a=T.bool(True), trans_b=T.bool(False))', # noqa: E501
+ 'd: T.handle = T.handle()\nT.cuda.tcgen05_encode_instr_descriptor(d,
K=16, M=16, N=16, a_dtype="f16", b_dtype="f16", d_dtype="f16", trans_a=True,
trans_b=False)', # noqa: E501
)
_assert_namespace_print(
cuda_op.cuda_tcgen05_encode_instr_descriptor_block_scaled(
@@ -795,8 +795,6 @@ def test_printer_ptx_more():
b_dtype="f16",
sfa_dtype="f16",
sfb_dtype="f16",
- sfa_tmem_addr=a,
- sfb_tmem_addr=b,
M=16,
N=16,
K=16,
@@ -808,9 +806,7 @@ def test_printer_ptx_more():
neg_b=False,
),
"d: T.handle = T.handle()\n"
- "a: T.handle = T.handle()\n"
- "b: T.handle = T.handle()\n"
- 'T.cuda.tcgen05.encode_instr_descriptor_block_scaled(d, d_dtype="f16",
a_dtype="f16", b_dtype="f16", sfa_dtype="f16", sfb_dtype="f16",
sfa_tmem_addr=a, sfb_tmem_addr=b, M=16, N=16, K=16, trans_a=T.bool(True),
trans_b=T.bool(False), is_sparse=T.bool(True))', # noqa: E501
+ 'T.cuda.tcgen05_encode_instr_descriptor_block_scaled(d, K=16, M=16,
N=16, a_dtype="f16", b_dtype="f16", d_dtype="f16", is_sparse=True,
sfa_dtype="f16", sfb_dtype="f16", trans_a=True, trans_b=False)', # noqa: E501
)
